fa3: sync-free eagle spec via fixed-window draft-extend metadata (#31364)
This commit is contained in:
@@ -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),
|
||||
|
||||
Reference in New Issue
Block a user