fix: Always set logits_all = True when using speculative decoding

abetlen · abetlen · commit cb791716b42e · 2024-02-12T16:19:05.000-05:00
diff --git a/llama_cpp/llama.py b/llama_cpp/llama.py
@@ -281,7 +281,7 @@ def __init__(
         )
         self.context_params.yarn_orig_ctx = yarn_orig_ctx if yarn_orig_ctx != 0 else 0
         self.context_params.mul_mat_q = mul_mat_q
-        self.context_params.logits_all = logits_all
+        self.context_params.logits_all = logits_all if draft_model is None else True # Must be set to True for speculative decoding
         self.context_params.embedding = embedding
         self.context_params.offload_kqv = offload_kqv
 

Original file line number	Diff line number	Diff line change
`@@ -281,7 +281,7 @@ def __init__(`
`281`	`281`	`)`
`282`	`282`	`self.context_params.yarn_orig_ctx = yarn_orig_ctx if yarn_orig_ctx != 0 else 0`
`283`	`283`	`self.context_params.mul_mat_q = mul_mat_q`
`284`		`- self.context_params.logits_all = logits_all`
	`284`	`+ self.context_params.logits_all = logits_all if draft_model is None else True # Must be set to True for speculative decoding`
`285`	`285`	`self.context_params.embedding = embedding`
`286`	`286`	`self.context_params.offload_kqv = offload_kqv`
`287`	`287`