diff --git a/.github/workflows/android.yml b/.github/workflows/android.yml
index 4ab49e1..ace796e 100644
--- a/.github/workflows/android.yml
+++ b/.github/workflows/android.yml
@@ -23,15 +23,20 @@ jobs:
           distribution: temurin
           java-version: 17
       # Clone Godot
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
         with:
           repository: godotengine/godot
 
       # Clone our module under the correct directory
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
         with:
           path: modules/rwkv
 
+      - uses: actions/checkout@v4
+        with:
+          repository: harrisonvanderbyl/treegen
+          path: modules/treegen
+
         # init submodules in the rwkv module
       - name: Init submodules
         run: |
diff --git a/.github/workflows/linux.yml b/.github/workflows/linux.yml
index 86299a7..6607f05 100644
--- a/.github/workflows/linux.yml
+++ b/.github/workflows/linux.yml
@@ -37,6 +37,10 @@ jobs:
         with:
           path: modules/rwkv
 
+      - uses: actions/checkout@v4
+        with:
+          repository: harrisonvanderbyl/treegen
+          path: modules/treegen
         # init submodules in the rwkv module
       - name: Init submodules
         run: |
diff --git a/.github/workflows/mac.yml b/.github/workflows/mac.yml
index 9c76d26..eed82f8 100644
--- a/.github/workflows/mac.yml
+++ b/.github/workflows/mac.yml
@@ -40,6 +40,11 @@ jobs:
         with:
           path: modules/rwkv
 
+      - uses: actions/checkout@v4
+        with:
+          repository: harrisonvanderbyl/treegen
+          path: modules/treegen
+
         # init submodules in the rwkv module
       - name: Init submodules
         run: |
diff --git a/.github/workflows/windows.yml b/.github/workflows/windows.yml
index d8a8e61..c0da308 100644
--- a/.github/workflows/windows.yml
+++ b/.github/workflows/windows.yml
@@ -31,12 +31,16 @@ jobs:
 
     steps:
       # Clone Godot
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
         with:
           repository: godotengine/godot
 
+      - uses: actions/checkout@v4
+        with:
+          repository: harrisonvanderbyl/treegen
+          path: modules/treegen
       # Clone our module under the correct directory
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
         with:
           path: modules/rwkv
 
diff --git a/rwkv.cuh b/rwkv.cuh
index cb06497..c708401 160000
--- a/rwkv.cuh
+++ b/rwkv.cuh
@@ -1 +1 @@
-Subproject commit cb06497e9ed49ad19e0f0499af62e349b1e62e45
+Subproject commit c708401d7e557c937f9359a371f66af73fb301c0