Fix "Avoid computing lse in Ragged Prefill when there's no prefix match" (#5555)

2025-05-05 10:32:17 -07:00
parent b8559764f6
commit 22da3d978f
3 changed files with 23 additions and 15 deletions
@@ -355,7 +355,7 @@ class FlashInferMLAAttnBackend(AttentionBackend):

        if self.forward_metadata.use_ragged:
            # ragged prefill
-            o, _ = self.prefill_wrapper_ragged.forward_return_lse(
+            o = self.prefill_wrapper_ragged.forward(
                qall,
                k.view(-1, layer.tp_k_head_num, layer.head_dim),
                v.view(-1, layer.tp_k_head_num, layer.v_head_dim),