[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
|
self.eagle_worker.chain_mtp_hidden_states
|
||||||
and ret.hidden_states is not None
|
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 = (
|
select_index = (
|
||||||
torch.arange(bs, device=self.model_runner.device)
|
torch.arange(bs, device=self.model_runner.device)
|
||||||
|
|||||||
@@ -506,6 +506,15 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker):
|
|||||||
dim=-1,
|
dim=-1,
|
||||||
)
|
)
|
||||||
ret_topk_p, ret_topk_index = fast_topk(probs, self.topk, 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:
|
if forward_batch.extend_seq_lens is not None:
|
||||||
rotate_input_ids_triton(
|
rotate_input_ids_triton(
|
||||||
forward_batch.input_ids,
|
forward_batch.input_ids,
|
||||||
|
|||||||
Reference in New Issue
Block a user