Skip to content
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion flashinfer/cascade.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
Loading