Skip to content

Commit

Permalink
bugfix: fix the python api of prefill wrapper + custom mask (flashinf…
Browse files Browse the repository at this point in the history
…er-ai#460)

some tests are failed because we use the wrong function for prefill
wrapper with custom mask, this pr fixes the issue.
  • Loading branch information
yzh119 authored Aug 21, 2024
1 parent 7c397cb commit 048560d
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion python/flashinfer/prefill.py
Original file line number Diff line number Diff line change
Expand Up @@ -1056,7 +1056,7 @@ def forward_return_lse(
True, # return LSE
)
else:
out, lse = self._wrapper.forward(
out, lse = self._wrapper.forward_custom_mask(
q,
self._qo_indptr_buf,
*_unpack_paged_kv_cache(paged_kv_cache, self._kv_layout),
Expand Down

0 comments on commit 048560d

Please sign in to comment.