diff --git a/python/pyproject.toml b/python/pyproject.toml
index a8e086c36..07011e304 100755
--- a/python/pyproject.toml
+++ b/python/pyproject.toml
@@ -66,7 +66,7 @@ dependencies = [
   "tiktoken",
   "tilelang==0.1.8",
   "timm==1.0.16",
-  "tokenspeed_mla==0.1.1",
+  "tokenspeed_mla==0.1.6",
   "torch_memory_saver>=0.0.9.post1",
   "torch==2.11.0",
   "torchao==0.17.0",