[Fix] Resolve TVM Compatibility Issue by Removing Unsupported Argument (#3159)

akshatsehgal · web-flow · commit 7175da511ada · 2025-03-10T00:09:52.000-04:00
This PR resolves a compatibility issue with the TVM library by
removing an unsupported argument. The issue arose when the
`TIRPagedKVCache` class was instantiated with an
`enable_disaggregation` argument, which is not recognized
by the TVM library.
diff --git a/python/mlc_llm/compiler_pass/dispatch_kv_cache_creation.py b/python/mlc_llm/compiler_pass/dispatch_kv_cache_creation.py
@@ -152,6 +152,9 @@ def create_tir_paged_kv_cache(
             "support_sliding_window_", relax.ShapeStructInfo([kwargs["support_sliding_window"]])
         )
 
+        # Remove the 'enable_disaggregation' argument
+        kwargs.pop("enable_disaggregation", None)
+
         with bb.function(
             name="create_tir_paged_kv_cache",
             params=[

Original file line number	Diff line number	Diff line change
`@@ -152,6 +152,9 @@ def create_tir_paged_kv_cache(`
`152`	`152`	`"support_sliding_window_", relax.ShapeStructInfo([kwargs["support_sliding_window"]])`
`153`	`153`	`)`
`154`	`154`
	`155`	`+ # Remove the 'enable_disaggregation' argument`
	`156`	`+ kwargs.pop("enable_disaggregation", None)`
	`157`	`+`
`155`	`158`	`with bb.function(`
`156`	`159`	`name="create_tir_paged_kv_cache",`
`157`	`160`	`params=[`