From ec6fcb93cb8b100e26db2bcf6a43613e9b05c7d1 Mon Sep 17 00:00:00 2001 From: Qiaolin Yu Date: Sun, 24 May 2026 21:36:16 -0700 Subject: [PATCH] [perf][spec decoding] Skip common_template in TRTLLMMLAMultiStepDraftBackend init (#26241) --- .../layers/attention/trtllm_mla_backend.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index 695e739f1..4a5fed11c 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mla_backend.py @@ -1235,3 +1235,22 @@ class TRTLLMMLAMultiStepDraftBackend(FlashInferMLAMultiStepDraftBackend): kv_indptr_buf=self.kv_indptr[i], 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, + )