From 44f4ea917c62840c3e608a6ec520a0c322760537 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Thu, 16 Jul 2026 12:31:50 -0700 Subject: [PATCH] fa3: sync-free eagle spec via fixed-window draft-extend metadata (#31364) --- .../layers/attention/flashattention_backend.py | 17 +++++++++++++---- .../runner/decode_cuda_graph_runner.py | 6 +++++- .../eagle_draft_extend_cuda_graph_runner.py | 5 ++++- 3 files changed, 22 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index fe22e6296..1230781be 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -824,10 +824,15 @@ class FlashAttentionBackend(AttentionBackend): forward_batch.req_pool_indices, : metadata.max_seq_len_k ] - if ( - any(forward_batch.extend_prefix_lens_cpu) - or forward_batch.forward_mode.is_draft_extend_v2() - ): + if forward_batch.forward_mode.is_draft_extend_v2(): + # Fixed-q window: the host max is a config constant, and + # extend_seq_lens_cpu may be None on the GPU-only spec path. + extend_seq_lens = forward_batch.extend_seq_lens + metadata.max_seq_len_q = self.speculative_num_draft_tokens + metadata.cu_seqlens_q = torch.nn.functional.pad( + torch.cumsum(extend_seq_lens, dim=0, dtype=torch.int32), (1, 0) + ) + elif any(forward_batch.extend_prefix_lens_cpu): extend_seq_lens = forward_batch.extend_seq_lens metadata.max_seq_len_q = max(forward_batch.extend_seq_lens_cpu) metadata.cu_seqlens_q = torch.nn.functional.pad( @@ -3022,6 +3027,10 @@ class FlashAttentionBackend(AttentionBackend): class FlashAttentionMultiStepBackend: + # Read by decide_needs_cpu_seq_lens (a missing flag defaults to True); + # the multi-step draft and draft-extend paths are device-side. + needs_cpu_seq_lens: bool = False + def __init__( self, model_runner: ModelRunner, diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index 1cde0bfd1..e9c91b0d0 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py @@ -162,7 +162,11 @@ def build_replay_fb_view( if forward_batch.seq_lens_sum is None else forward_batch.seq_lens_sum + (bs - raw_bs) * seq_len_fill_value ), - seq_lens_cpu=buffers.seq_lens_cpu[:bs], + # Propagate mirror absence: the pinned buffer is not refreshed when the + # batch has no CPU mirror; a stale non-None tensor defeats None-guards. + seq_lens_cpu=( + None if forward_batch.seq_lens_cpu is None else buffers.seq_lens_cpu[:bs] + ), num_padding=bs - raw_bs, encoder_lens=buffers.encoder_lens[:bs] if is_encoder_decoder else None, out_cache_loc=getattr(forward_batch, "out_cache_loc", None), diff --git a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py index 1234aece3..71844c60d 100644 --- a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py @@ -584,7 +584,10 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): req_pool_indices=buffers.req_pool_indices, seq_lens=buffers.seq_lens, seq_lens_sum=seq_lens_sum, - seq_lens_cpu=buffers.seq_lens_cpu, + # Mirror absence must survive replay (stale buffer defeats None-guards). + seq_lens_cpu=( + None if forward_batch.seq_lens_cpu is None else buffers.seq_lens_cpu + ), encoder_lens=None, out_cache_loc=buffers.out_cache_loc[:num_tokens], out_cache_loc_dsv4=getattr(forward_batch, "out_cache_loc_dsv4", None),