[sgl] potential chained spec v2 fixes (#22041)
Co-authored-by: Mook <Godmook@users.noreply.github.com> Co-authored-by: yudian0504 <yudian0504@users.noreply.github.com>
This commit is contained in:
co-authored by
Mook
yudian0504
parent
7f2fcc0b08
commit
8f337682bd
@@ -437,7 +437,7 @@ class MultiLayerEagleDraftExtendCudaGraphRunner:
|
||||
self.eagle_worker.chain_mtp_hidden_states
|
||||
and ret.hidden_states is not None
|
||||
):
|
||||
self.hidden_states[:num_tokens].copy_(ret.hidden_states[:num_tokens])
|
||||
buffers.hidden_states[:num_tokens].copy_(ret.hidden_states[:num_tokens])
|
||||
|
||||
select_index = (
|
||||
torch.arange(bs, device=self.model_runner.device)
|
||||
|
||||
@@ -506,6 +506,15 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker):
|
||||
dim=-1,
|
||||
)
|
||||
ret_topk_p, ret_topk_index = fast_topk(probs, self.topk, dim=-1)
|
||||
# Chain-style: use this step's output hidden_states as next step's input
|
||||
if (
|
||||
self.chain_mtp_hidden_states
|
||||
and step < self.speculative_num_steps - 1
|
||||
and draft_logits_output.logits_output.hidden_states is not None
|
||||
):
|
||||
forward_batch.spec_info.hidden_states = (
|
||||
draft_logits_output.logits_output.hidden_states
|
||||
)
|
||||
if forward_batch.extend_seq_lens is not None:
|
||||
rotate_input_ids_triton(
|
||||
forward_batch.input_ids,
|
||||
|
||||
Reference in New Issue
Block a user