[spec] Fix index_share_for_mtp_iteration being a no-op in EAGLE MTP draft (#29654)
This commit is contained in:
@@ -430,8 +430,8 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
|
|||||||
# For hidden states before normal
|
# For hidden states before normal
|
||||||
return_hidden_states_before_norm: bool = False
|
return_hidden_states_before_norm: bool = False
|
||||||
|
|
||||||
# For NSA/DSA topk_indices reuse across forward calls (e.g., EAGLE draft)
|
# Gate for reusing the first MTP draft step's indexer topk across steps;
|
||||||
topk_indices: Optional[torch.Tensor] = None
|
# the carried topk lives on spec_info (see EagleDraftInput.mtp_topk_indices).
|
||||||
reuse_mtp_topk_indices: Optional[bool] = False
|
reuse_mtp_topk_indices: Optional[bool] = False
|
||||||
|
|
||||||
# === Forward-derived (built in init_new on the forward stream; FB-owned) ===
|
# === Forward-derived (built in init_new on the forward stream; FB-owned) ===
|
||||||
|
|||||||
@@ -228,13 +228,13 @@ class DeepseekModelNextN(nn.Module):
|
|||||||
residual,
|
residual,
|
||||||
zero_allocator,
|
zero_allocator,
|
||||||
prev_topk_indices=(
|
prev_topk_indices=(
|
||||||
forward_batch.topk_indices
|
forward_batch.spec_info.mtp_topk_indices
|
||||||
if forward_batch.reuse_mtp_topk_indices
|
if forward_batch.reuse_mtp_topk_indices
|
||||||
else None
|
else None
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
if forward_batch.reuse_mtp_topk_indices:
|
if forward_batch.reuse_mtp_topk_indices:
|
||||||
forward_batch.topk_indices = topk_indices
|
forward_batch.spec_info.mtp_topk_indices = topk_indices
|
||||||
|
|
||||||
if not forward_batch.forward_mode.is_idle():
|
if not forward_batch.forward_mode.is_idle():
|
||||||
if residual is not None:
|
if residual is not None:
|
||||||
|
|||||||
@@ -159,6 +159,10 @@ class EagleDraftInput(SpecInput):
|
|||||||
hidden_states: Optional[torch.Tensor] = None
|
hidden_states: Optional[torch.Tensor] = None
|
||||||
capture_hidden_mode: CaptureHiddenMode = CaptureHiddenMode.FULL
|
capture_hidden_mode: CaptureHiddenMode = CaptureHiddenMode.FULL
|
||||||
|
|
||||||
|
# Survives across draft steps: spec_info is shared by reference across the
|
||||||
|
# per-step forwards (each runs on a copied ForwardBatch, dropping writebacks).
|
||||||
|
mtp_topk_indices: Optional[torch.Tensor] = None
|
||||||
|
|
||||||
# Per-req bonus token (the "+1" target prediction at end of each accept
|
# Per-req bonus token (the "+1" target prediction at end of each accept
|
||||||
# chain); the worker copies it here post-extend for next iter's draft.
|
# chain); the worker copies it here post-extend for next iter's draft.
|
||||||
bonus_tokens: torch.Tensor = None
|
bonus_tokens: torch.Tensor = None
|
||||||
|
|||||||
@@ -612,7 +612,7 @@ class EagleDraftWorker(EagleDraftWorkerBase):
|
|||||||
scores = None
|
scores = None
|
||||||
if self.index_share_for_mtp_iteration:
|
if self.index_share_for_mtp_iteration:
|
||||||
forward_batch.reuse_mtp_topk_indices = True
|
forward_batch.reuse_mtp_topk_indices = True
|
||||||
forward_batch.topk_indices = None
|
spec_info.mtp_topk_indices = None
|
||||||
for i in range(self.speculative_num_steps):
|
for i in range(self.speculative_num_steps):
|
||||||
input_ids, hidden_states, scores, tree_info = select_top_k_tokens(
|
input_ids, hidden_states, scores, tree_info = select_top_k_tokens(
|
||||||
i, topk_p, topk_index, hidden_states, scores, self.topk
|
i, topk_p, topk_index, hidden_states, scores, self.topk
|
||||||
@@ -688,7 +688,7 @@ class EagleDraftWorker(EagleDraftWorkerBase):
|
|||||||
forward_batch.positions.add_(1)
|
forward_batch.positions.add_(1)
|
||||||
|
|
||||||
if self.index_share_for_mtp_iteration:
|
if self.index_share_for_mtp_iteration:
|
||||||
forward_batch.topk_indices = None
|
spec_info.mtp_topk_indices = None
|
||||||
forward_batch.reuse_mtp_topk_indices = False
|
forward_batch.reuse_mtp_topk_indices = False
|
||||||
|
|
||||||
# Organize the results
|
# Organize the results
|
||||||
|
|||||||
Reference in New Issue
Block a user