bugfix: fix the python api of prefill wrapper + custom mask (#460)

some tests are failed because we use the wrong function for prefill wrapper with custom mask, this pr fixes the issue.
flashinfer-ai · Aug 21, 2024 · 048560d · 048560d
1 parent 7c397cb
commit 048560d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/flashinfer/prefill.py b/python/flashinfer/prefill.py
@@ -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),