[Fix] Resolve performance drop in speculative decoding aiter backend (#11087)

59911195 · jacky.cheng · GitHub · 424591d5 · 59911195
Unverified Commit 59911195 authored Sep 30, 2025 by jacky.cheng Committed by GitHub Sep 29, 2025
Hide whitespace changes
Inline Side-by-side

Showing with 5 additions and 1 deletion

python/sglang/srt/layers/attention/aiter_backend.py python/sglang/srt/layers/attention/aiter_backend.py +5 -1

No files found.
--- a/python/sglang/srt/layers/attention/aiter_backend.py
+++ b/python/sglang/srt/layers/attention/aiter_backend.py
@@ -619,7 +619,11 @@ class AiterAttnBackend(AttentionBackend):
            assert len(k.shape) == 3
            assert len(v.shape) == 3

-            if forward_batch.forward_mode.is_extend():
+            if (
+                forward_batch.forward_mode.is_extend()
+                and not forward_batch.forward_mode.is_target_verify()
+                and not forward_batch.forward_mode.is_draft_extend()
+            ):
                if kv_indices.shape[0] == 0:
                    o = flash_attn_varlen_func(
                        q,