Skip to content

Commit d7f0b0a

Browse files
authored
[fix] Correct the returned value of has_spec_drafter (#6178)
Signed-off-by: ziyixiong-nv <[email protected]>
1 parent f9b0a91 commit d7f0b0a

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

tensorrt_llm/_torch/speculative/interface.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,8 @@ def has_spec_decoder(self):
7777
return self.is_mtp() or self.is_eagle3() or self.is_eagle3_one_model()
7878

7979
def has_spec_drafter(self):
80-
return self.is_ngram() or self.is_user_provided()
80+
return self.is_eagle3() or self.is_draft_target() or self.is_ngram(
81+
) or self.is_user_provided()
8182

8283
def extend_ctx(self, attention_backend: Type[AttentionBackend]):
8384
"""

0 commit comments

Comments
 (0)