diff --git a/flashinfer/cascade.py b/flashinfer/cascade.py
index ff562d8d9b..011214a491 100644
--- a/flashinfer/cascade.py
+++ b/flashinfer/cascade.py
@@ -544,7 +544,10 @@ def run(
         )
         for wrapper in self._batch_prefill_wrappers[:-1]:
             if _HIP_FUSED_CASCADE:
-                out, lse = wrapper.run(
+                # mypy sees prefill.py (CUDA) wrapper here; at runtime on HIP,
+                # flashinfer.__init__ aliases sys.modules["flashinfer.prefill"]
+                # to prefill_rocm, whose run() accepts partial_state.
+                out, lse = wrapper.run(  # type: ignore[call-overload]
                     q, paged_kv_cache, return_lse=True, partial_state=(out, lse)
                 )
             else: