diff --git a/backends/exllamav3/model.py b/backends/exllamav3/model.py index c1d5074d..df0b1bd2 100644 --- a/backends/exllamav3/model.py +++ b/backends/exllamav3/model.py @@ -259,6 +259,7 @@ class ExllamaV3Container: max_batch_size: Optional[int] = None draft_num_tokens: Optional[int] = None dynamic_draft: Optional[bool] = False + draft_confidence: float = 0.4 ngram_match_min: int = 0 hybrid_draft_tokens: int = 0 recurrent_checkpoint_interval: Optional[int] = None @@ -419,6 +420,9 @@ async def create(cls, model_directory: pathlib.Path, hf_model: HFModel, **kwargs else None ) self.dynamic_draft = draft_args.get("dynamic_draft", False) + self.draft_confidence = unwrap(draft_args.get("draft_confidence"), 0.4) + if not 0.0 < self.draft_confidence < 1.0: + raise ValueError("draft_confidence must be between 0 and 1") # Long suffix-match drafting on top of whichever drafting mode is # active, or on its own: a job repeating its own context verbatim gets @@ -1222,6 +1226,7 @@ async def create_generator(self): cpu_cache_size=config.memory.sysmem_kv_cache * 1024**2, num_draft_tokens=self.draft_num_tokens, dynamic_draft_tokens=self.dynamic_draft, + draft_confidence=self.draft_confidence, ngram_match_min=self.ngram_match_min, hybrid_draft_tokens=self.hybrid_draft_tokens, **checkpoint_kwargs, diff --git a/common/config_models.py b/common/config_models.py index 4aa37dec..eff94ee9 100644 --- a/common/config_models.py +++ b/common/config_models.py @@ -678,6 +678,19 @@ class DraftModelConfig(BaseConfigModel): "Ceiling is given by num_draft_tokens." ), ) + draft_confidence: Optional[float] = Field( + 0.4, + description=( + "Target acceptance probability for drafted positions when dynamic_draft is " + "enabled (default: 0.4).\n" + "Evaluated against an online mapping from drafter confidence scores to " + "observed acceptance rates.\n" + "Lower values keep longer drafts; higher values truncate more aggressively.\n" + "Must be between 0 and 1. Ignored for n-gram drafting." + ), + gt=0, + lt=1, + ) ngram_match_min: Optional[int] = Field( 2, description=( diff --git a/config_sample.yml b/config_sample.yml index 795f6593..d94ca2bb 100644 --- a/config_sample.yml +++ b/config_sample.yml @@ -395,6 +395,12 @@ draft_model: # Ceiling is given by num_draft_tokens. dynamic_draft: + # Target acceptance probability for drafted positions when dynamic_draft is enabled + # (default: 0.4). Evaluated against an online mapping from drafter confidence scores + # to observed acceptance rates. Lower values keep longer drafts; higher values + # truncate more aggressively. Must be between 0 and 1. Ignored for n-gram drafting. + draft_confidence: 0.4 + # Minimum match length for exllamav3 n-gram drafting (default: 2). # Only used when draft_mode is ngram. ngram_match_min: 2 diff --git a/docs/02.-Server-options.md b/docs/02.-Server-options.md index a17b2cf6..825af559 100644 --- a/docs/02.-Server-options.md +++ b/docs/02.-Server-options.md @@ -126,6 +126,7 @@ Note: Sub-block of Model Options. Same rules apply. | draft_gpu_split | List[Float] ([]) | Float array of GBs to split a draft model between GPUs. | | draft_num_tokens | Integer (None) | Number of tokens to draft per iteration. Exllamav3 selects the mode-specific default when blank. | | dynamic_draft | Bool (False) | Adjust the number of draft tokens dynamically based on observed acceptance rates. Ceiling is given by `draft_num_tokens`. | +| draft_confidence | Float (0.4) | Target acceptance probability for drafted positions when `dynamic_draft` is enabled. Evaluated against an online mapping from drafter confidence scores to observed acceptance rates. Lower values keep longer drafts; higher values truncate more aggressively. Between 0 and 1. Ignored for n-gram drafting. | | ngram_match_min | Integer (2) | Minimum match length for exllamav3 n-gram drafting. Only used by ngram mode. | | hybrid_draft_tokens | Integer (0) | Hybrid drafting: when a response is repeating a passage of its context verbatim (reciting a drafted reply, editing source), the earlier continuation is drafted this many tokens at a time and verified in one pass. Works on top of any draft mode, or on its own; 0 disables. Longer values advance more tokens at prefill speed on a hit, at the cost of a wider verification pass. | diff --git a/endpoints/core/types/model.py b/endpoints/core/types/model.py index 62cf23e4..09012a16 100644 --- a/endpoints/core/types/model.py +++ b/endpoints/core/types/model.py @@ -102,6 +102,13 @@ class DraftModelLoadRequest(_WarnOnUnknownFields): draft_cache_mode: Optional[str] = None draft_num_tokens: Optional[int] = None dynamic_draft: Optional[bool] = None + draft_confidence: Optional[float] = Field( + description="Target acceptance probability for dynamic drafting, in (0, 1)", + default=None, + examples=[0.4], + gt=0, + lt=1, + ) ngram_match_min: Optional[int] = None hybrid_draft_tokens: Optional[int] = None draft_rope_scale: Optional[float] = None