Fix stale flashinfer-MLA fallback poisoning spec verify capture (trtllm_mla + tc_piecewise) (#32288)

python/sglang/srt/layers/attention/trtllm_mla_backend.py

@@ -857,8 +857,17 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
         llama_4_scaling: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
 
+        # The fallback belongs to genuine extend forwards only. Target-verify /
+        # draft-extend must never honor it: `forward_prefill_metadata` is a
+        # stale leftover from the last prefill there (eager init clears it,
+        # but decode-graph capture does not), and capturing verify through the
+        # flashinfer path binds the graph to prefill-planned wrapper buffers,
+        # which fault (illegal address) at replay.
         if (
-            self.forward_prefill_metadata is not None
+            forward_batch.forward_mode.is_extend()
+            and not forward_batch.forward_mode.is_target_verify()
+            and not forward_batch.forward_mode.is_draft_extend_v2()
+            and self.forward_prefill_metadata is not None
             and self.forward_prefill_metadata.fallback_to_flashinfer_impl
         ):
             return super().forward_extend(

python/sglang/srt/server_args.py

@@ -4225,19 +4225,6 @@ class ServerArgs:
                 "decode context parallel (dcp_size > 1)",
                 lambda: self.dcp_size > 1,
             ),
-            # TcPiecewise makes the trtllm_mla prefill fall back to the
-            # flashinfer-MLA implementation, which faults (illegal address)
-            # on an FP8 KV cache.
-            (
-                "MLA attention with FP8 KV cache",
-                lambda: self.kv_cache_dtype.startswith("fp8")
-                and (
-                    _resolved_view(self).attention_backend
-                    in ("trtllm_mla", "flashinfer_mla")
-                    or _resolved_view(self).prefill_attention_backend
-                    in ("trtllm_mla", "flashinfer_mla")
-                ),
-            ),
         ]
         for _name, predicate in rules:
             if predicate():