Reduce one step decode for draft model. (#11561)
This commit is contained in:
@@ -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]
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user