Unverified Commit 6d036468 authored by Shangming Cai's avatar Shangming Cai Committed by GitHub
Browse files

Fix 1-step draft model forward (#11653)


Signed-off-by: default avatarShangming Cai <csmthu@gmail.com>
Co-authored-by: default avatarLiangsheng Yin <lsyincs@gmail.com>
parent 8221f9ae
...@@ -33,15 +33,7 @@ class DraftBackendFactory: ...@@ -33,15 +33,7 @@ class DraftBackendFactory:
def create_decode_backend(self): def create_decode_backend(self):
if self.speculative_num_steps == 1: if self.speculative_num_steps == 1:
return None
class DummyAttnBackend:
def __init__(self):
pass
def init_forward_metadata(*args, **kwargs):
pass
return DummyAttnBackend()
backend_map = { backend_map = {
"flashinfer": self._create_flashinfer_decode_backend, "flashinfer": self._create_flashinfer_decode_backend,
......
...@@ -276,7 +276,7 @@ class EagleVerifyInputV2Mixin: ...@@ -276,7 +276,7 @@ class EagleVerifyInputV2Mixin:
accept_length=accept_length, # mutable accept_length=accept_length, # mutable
simulate_acc_len=SIMULATE_ACC_LEN, simulate_acc_len=SIMULATE_ACC_LEN,
bs=bs, bs=bs,
spec_steps=self.draft_token_num, spec_steps=self.spec_steps,
) )
# Include the bonus token # Include the bonus token
......
...@@ -218,6 +218,7 @@ class EAGLEWorker(TpModelWorker): ...@@ -218,6 +218,7 @@ class EAGLEWorker(TpModelWorker):
return return
# Capture draft # Capture draft
if self.speculative_num_steps > 1:
tic = time.perf_counter() tic = time.perf_counter()
before_mem = get_available_gpu_memory(self.device, self.gpu_id) before_mem = get_available_gpu_memory(self.device, self.gpu_id)
logger.info( logger.info(
...@@ -500,8 +501,11 @@ class EAGLEWorker(TpModelWorker): ...@@ -500,8 +501,11 @@ class EAGLEWorker(TpModelWorker):
) )
else: else:
forward_batch.can_run_dp_cuda_graph = False forward_batch.can_run_dp_cuda_graph = False
if not forward_batch.forward_mode.is_idle(): if (
# Initialize attention backend not forward_batch.forward_mode.is_idle()
and self.speculative_num_steps > 1
):
# Skip attention backend init for idle mode or 1-step draft
self.draft_attn_backend.init_forward_metadata(forward_batch) self.draft_attn_backend.init_forward_metadata(forward_batch)
# Run forward steps # Run forward steps
parent_list, top_scores_index, draft_tokens = self.draft_forward( parent_list, top_scores_index, draft_tokens = self.draft_forward(
......
...@@ -97,6 +97,9 @@ class EAGLEWorkerV2(EAGLEWorker): ...@@ -97,6 +97,9 @@ class EAGLEWorkerV2(EAGLEWorker):
forward_batch, forward_batch,
) )
else: else:
if self.speculative_num_steps > 1:
# Skip attention backend init for 1-step draft,
# `draft_forward` only does sample in this case.
self.draft_attn_backend.init_forward_metadata(forward_batch) self.draft_attn_backend.init_forward_metadata(forward_batch)
parent_list, top_scores_index, draft_tokens = self.draft_forward( parent_list, top_scores_index, draft_tokens = self.draft_forward(
forward_batch forward_batch
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment