[Bugfix][Spec Decode] Drop FlashAttention's AOT schedule for a sliding-window DFlash drafter (#54374)
Zihan Zhang committed
Sep 4, 2026 at 13:42 UTC
5093e4844a75f279d1921d7e34ccd23d68529f27
1 file changed
+17
-1
vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
+17
-1
@@ -15,7 +15,10 @@ from vllm.logger import init_logger
15
from vllm.triton_utils import tl, triton
16
from vllm.v1.attention.backend import AttentionCGSupport
17
from vllm.v1.attention.backends.utils import PAD_SLOT_ID
18
-from vllm.v1.kv_cache_interface import KVCacheConfig
18
+from vllm.v1.kv_cache_interface import (
19
+ KVCacheConfig,
20
+ get_kv_cache_spec_sliding_window,
21
+)
22
from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
23
from vllm.v1.worker.gpu.block_table import BlockTables
24
from vllm.v1.worker.gpu.cp_utils import cp_local_slot
@@ -178,6 +181,19 @@ class DFlashSpeculator(DraftModelSpeculator):
181
target_attn_groups,
182
)
183
184
+ # FlashAttention's AOT split schedule is wrong for a windowed drafter,
185
+ # and `_get_sliding_window_configs` leaves it on or off depending on
186
+ # whether the target also runs FlashAttention. Decide it here instead.
187
+ for groups in self.attn_groups:
188
+ for group in groups:
189
+ builder = group.get_metadata_builder()
190
+ if getattr(
191
+ builder, "aot_schedule", False
192
+ ) and get_kv_cache_spec_sliding_window(builder.kv_cache_spec):
193
+ # `aot_schedule` belongs to FlashAttention's builder, not
194
+ # to the base class this loop is typed against.
195
+ builder.aot_schedule = False # type: ignore[attr-defined]
196
+
197
self.draft_kv_cache_group_ids = [
198
gid for gid, g in enumerate(self.attn_groups) if g
199
]