diff --git a/BMM/BMM_Theory/bmm_theory/branches/__init__.py b/BMM/BMM_Theory/bmm_theory/branches/__init__.py
index 2c955b2..3894e03 100644
--- a/BMM/BMM_Theory/bmm_theory/branches/__init__.py
+++ b/BMM/BMM_Theory/bmm_theory/branches/__init__.py
@@ -1,4 +1,11 @@
 from .merge_batch import MergeBatchBranch
 from .iter_batch import IterBatchBranch
+from .to_matmul import ToMatmulBranch
+from .special import SpecialBranch
+from .stream_k import StreamKBranch
+from .asw_basic import AswBasicBranch
 
-__all__ = ["MergeBatchBranch", "IterBatchBranch"]
+__all__ = [
+    "MergeBatchBranch", "IterBatchBranch", "ToMatmulBranch",
+    "SpecialBranch", "StreamKBranch", "AswBasicBranch",
+]