From 82fe0f041aec90b560fb39b8c1c2687b9dd0029e Mon Sep 17 00:00:00 2001 From: Qiaolin Yu Date: Fri, 24 Jul 2026 13:24:26 -0700 Subject: [PATCH] Fix stale flashinfer-MLA fallback poisoning spec verify capture (trtllm_mla + tc_piecewise) (#32288) --- .../srt/layers/attention/trtllm_mla_backend.py | 11 ++++++++++- python/sglang/srt/server_args.py | 13 ------------- 2 files changed, 10 insertions(+), 14 deletions(-) diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index 79a1766a3..37016365a 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/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( diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 715d79604..bf74ffc25 100644 --- a/python/sglang/srt/server_args.py +++ b/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():