[perf][spec decoding] Skip common_template in TRTLLMMLAMultiStepDraftBackend init (#26241)
This commit is contained in:
@@ -1235,3 +1235,22 @@ class TRTLLMMLAMultiStepDraftBackend(FlashInferMLAMultiStepDraftBackend):
|
|||||||
kv_indptr_buf=self.kv_indptr[i],
|
kv_indptr_buf=self.kv_indptr[i],
|
||||||
q_indptr_decode_buf=self.q_indptr_decode,
|
q_indptr_decode_buf=self.q_indptr_decode,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
def init_forward_metadata(self, forward_batch: ForwardBatch):
|
||||||
|
for i in range(self.speculative_num_steps - 1):
|
||||||
|
self.attn_backends[i].init_forward_metadata(forward_batch)
|
||||||
|
|
||||||
|
def init_forward_metadata_replay_cuda_graph(
|
||||||
|
self, forward_batch: ForwardBatch, bs: int
|
||||||
|
):
|
||||||
|
for i in range(self.speculative_num_steps - 1):
|
||||||
|
self.attn_backends[i].init_forward_metadata_replay_cuda_graph(
|
||||||
|
bs,
|
||||||
|
forward_batch.req_pool_indices,
|
||||||
|
forward_batch.seq_lens,
|
||||||
|
seq_lens_sum=None,
|
||||||
|
encoder_lens=None,
|
||||||
|
forward_mode=ForwardMode.DECODE,
|
||||||
|
spec_info=forward_batch.spec_info,
|
||||||
|
seq_lens_cpu=forward_batch.seq_lens_cpu,
|
||||||
|
)
|
||||||
|
|||||||
Reference in New Issue
Block a user