fa3: sync-free eagle spec via fixed-window draft-extend metadata (#31364)

This commit is contained in:
Liangsheng Yin
2026-07-16 12:31:50 -07:00
committed by GitHub
parent 1f34911de7
commit 44f4ea917c
3 changed files with 22 additions and 6 deletions
@@ -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,
@@ -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),
@@ -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),