Revert "Avoid computing lse in Ragged Prefill when there's no prefix.… (#5544)

2025-04-18 16:50:21 -07:00
parent 08b518d51f
commit a6f892e5d0
3 changed files with 12 additions and 19 deletions
@@ -348,7 +348,7 @@ class FlashInferMLAAttnBackend(AttentionBackend):

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