--- a/mlc_llm/nn/rnn_state.py
+++ b/mlc_llm/nn/rnn_state.py
@@ -1,3 +1,4 @@
+from tvm.script import ir as I
 """RNN State modeling."""
 
 from collections.abc import Sequence
@@ -194,8 +195,8 @@
                     for s in range(shape[0]):
                         with T.sblock("copy"):
                             vi, vs = T.axis.remap("SS", [i, s])
-                            seq_id: T.int32 = seq_slot_ids[vi]
-                            history_id: T.int32 = history_slot_ids[vi]
+                            seq_id = I.meta_var(seq_slot_ids[vi])
+                            history_id = I.meta_var(history_slot_ids[vi])
                             output[vi, vs] = storage[seq_id, history_id, vs]
 
             return f
@@ -221,8 +222,8 @@
                     for s in T.grid(*shape):
                         with T.sblock("copy"):
                             vi, *vs = T.axis.remap("S" * (len(shape) + 1), [i, *s])
-                            seq_id: T.int32 = seq_slot_ids[vi]
-                            history_id: T.int32 = history_slot_ids[vi]
+                            seq_id = I.meta_var(seq_slot_ids[vi])
+                            history_id = I.meta_var(history_slot_ids[vi])
                             # The following line is equivalent to:
                             # `output[vi, *vs] = storage[seq_id, history_id, *vs]`
                             # However, unpacking operator in subscript requires Python 3.11 or newer
@@ -291,10 +292,10 @@
                     for s in range(shape[0]):
                         with T.sblock("copy"):
                             vi, vs = T.axis.remap("SS", [i, s])
-                            seq_id: T.int32 = seq_slot_ids[vi]
-                            history_id: T.int32 = (history_slot_ids[vi] + 1) % T.cast(
+                            seq_id = I.meta_var(seq_slot_ids[vi])
+                            history_id = I.meta_var((history_slot_ids[vi] + 1) % T.cast(
                                 max_history, "int32"
-                            )
+                            ))
                             storage[seq_id, history_id, vs] = data[vi, vs]
 
             return f
@@ -319,10 +320,10 @@
                     for s in T.grid(*shape):
                         with T.sblock("copy"):
                             vi, *vs = T.axis.remap("S" * (len(shape) + 1), [i, *s])
-                            seq_id: T.int32 = seq_slot_ids[vi]
-                            history_id: T.int32 = (history_slot_ids[vi] + 1) % T.cast(
+                            seq_id = I.meta_var(seq_slot_ids[vi])
+                            history_id = I.meta_var((history_slot_ids[vi] + 1) % T.cast(
                                 max_history, "int32"
-                            )
+                            ))
                             # The following line is equivalent to:
                             # `storage[seq_id, history_id, *vs] = data[vi, *vs]`
                             # However, unpacking operator in subscript requires Python 3.11 or newer
--- a/mlc_llm/compiler_pass/low_batch_specialization.py
+++ b/mlc_llm/compiler_pass/low_batch_specialization.py
@@ -34,7 +34,7 @@
                     for low_batch_func in low_batch_funcs
                 ):
                     continue
-                buffers = [param for param in func.params if tirx.is_buffer_var(param)]
+                buffers = list(func.buffer_map.values())
                 shapes = [buffer.shape for buffer in buffers]
                 symbolic_vars = set(
                     expr for shape in shapes for expr in shape if isinstance(expr, tirx.Var)
--- a/mlc_llm/compiler_pass/lift_global_buffer_alloc.py
+++ b/mlc_llm/compiler_pass/lift_global_buffer_alloc.py
@@ -96,19 +96,22 @@
     """Remove the global buffer allocation for a given TIR PrimFunc."""
     assert isinstance(func.body, tirx.SBlockRealize)
     params = list(func.params)
+    buffer_map = dict(func.buffer_map)
     tensor_sinfo = []
     alloc_buffers = []
 
     insertion_point = len(params)
-    while not tirx.is_buffer_var(params[insertion_point - 1]):
+    while not isinstance(params[insertion_point - 1].ty, tvm.ir.PointerType):
         insertion_point -= 1
         assert insertion_point >= 1
 
     prev_root_block = func.body.block
     for buf_alloc in func.body.block.alloc_buffers:
         if buf_alloc.scope() == "global":
-            params.insert(insertion_point, buf_alloc)
+            param = tirx.Var("var_" + buf_alloc.name, "handle")
+            params.insert(insertion_point, param)
             insertion_point += 1
+            buffer_map[param] = buf_alloc
             tensor_sinfo.append(relax.TensorType(buf_alloc.shape, buf_alloc.dtype))
         else:
             alloc_buffers.append(buf_alloc)
@@ -136,6 +139,7 @@
         params=params,
         body=tirx.SBlockRealize(iter_values=[], predicate=True, block=root_block),
         ret_type=func.ret_type,
+        buffer_map=buffer_map,
         attrs=func.attrs,
     )
     return updated_func, tensor_sinfo
@@ -159,7 +163,7 @@
 
     n_arg = len(call.args[1].fields)
     for i in range(n_arg):
-        buffer_shape = func.params[i].shape
+        buffer_shape = func.buffer_map[func.params[i]].shape
         arg_shape = call.args[1][i].ty.shape.values
         assert len(buffer_shape) == len(arg_shape)
         for v_l, v_r in zip(buffer_shape, arg_shape):
@@ -173,7 +177,7 @@
         [ret_tensors] if isinstance(ret_tensors, relax.TensorType) else list(ret_tensors.fields)
     )
     for i, ret_tensor in enumerate(ret_tensors):
-        buffer_shape = func.params[n_arg + i].shape
+        buffer_shape = func.buffer_map[func.params[n_arg + i]].shape
         ret_tensor_shape = ret_tensor.shape.values
         assert len(buffer_shape) == len(ret_tensor_shape)
         for v_l, v_r in zip(buffer_shape, ret_tensor_shape):