Reduce one step decode for draft model. (#11561)

This commit is contained in:
Liangsheng Yin
2025-10-14 23:52:04 +08:00
committed by GitHub
parent 56222658ec
commit 5ea96ac7cc
10 changed files with 253 additions and 211 deletions
@@ -916,7 +916,7 @@ class FlashInferMLAMultiStepDraftBackend:
)
self.attn_backends = []
for i in range(self.speculative_num_steps):
for i in range(self.speculative_num_steps - 1):
self.attn_backends.append(
FlashInferMLAAttnBackend(
model_runner,
@@ -998,7 +998,7 @@ class FlashInferMLAMultiStepDraftBackend:
device="cuda",
)
for i in range(self.speculative_num_steps):
for i in range(self.speculative_num_steps - 1):
self.attn_backends[i].init_cuda_graph_state(
max_bs, max_num_tokens, kv_indices_buf=self.cuda_graph_kv_indices[i]
)