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 forward_batch.req_pool_indices, : metadata.max_seq_len_k
] ]
if ( if forward_batch.forward_mode.is_draft_extend_v2():
any(forward_batch.extend_prefix_lens_cpu) # Fixed-q window: the host max is a config constant, and
or forward_batch.forward_mode.is_draft_extend_v2() # 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 extend_seq_lens = forward_batch.extend_seq_lens
metadata.max_seq_len_q = max(forward_batch.extend_seq_lens_cpu) metadata.max_seq_len_q = max(forward_batch.extend_seq_lens_cpu)
metadata.cu_seqlens_q = torch.nn.functional.pad( metadata.cu_seqlens_q = torch.nn.functional.pad(
@@ -3022,6 +3027,10 @@ class FlashAttentionBackend(AttentionBackend):
class FlashAttentionMultiStepBackend: 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__( def __init__(
self, self,
model_runner: ModelRunner, model_runner: ModelRunner,
@@ -162,7 +162,11 @@ def build_replay_fb_view(
if forward_batch.seq_lens_sum is None if forward_batch.seq_lens_sum is None
else forward_batch.seq_lens_sum + (bs - raw_bs) * seq_len_fill_value 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, num_padding=bs - raw_bs,
encoder_lens=buffers.encoder_lens[:bs] if is_encoder_decoder else None, encoder_lens=buffers.encoder_lens[:bs] if is_encoder_decoder else None,
out_cache_loc=getattr(forward_batch, "out_cache_loc", 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, req_pool_indices=buffers.req_pool_indices,
seq_lens=buffers.seq_lens, seq_lens=buffers.seq_lens,
seq_lens_sum=seq_lens_sum, 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, encoder_lens=None,
out_cache_loc=buffers.out_cache_loc[:num_tokens], out_cache_loc=buffers.out_cache_loc[:num_tokens],
out_cache_loc_dsv4=getattr(forward_batch, "out_cache_loc_dsv4", None), out_cache_loc_dsv4=getattr(forward_batch, "out_cache_loc_dsv4", None),