fix(examples): honor DEEPINFER_USE_NEXTN_DECODE draft length

offline_inference.py and benchmark.py wrap the env value in min(1, ...),
which caps lookahead_tokens at 1 for any configured value above zero.
The README documents DEEPINFER_USE_NEXTN_DECODE as controlling the
maximum draft length (default: 5), so pass the configured value through.

Co-Authored-By: Claude Code <noreply@anthropic.com>
This commit is contained in:
uncle 2026-10-21 16:58:12 +08:00
parent 9bd9cf243f
commit 0e80d7c835
2 changed files with 2 additions and 2 deletions

View File

@ -181,7 +181,7 @@ def setup_llm(dp_rank: int) -> LLM:
max_num_batched_tokens=args.prefill_batch_size,
max_model_len=max_model_len,
gpu_memory_utilization=0.95,
lookahead_tokens=min(1, int(os.getenv("DEEPINFER_USE_NEXTN_DECODE", "5"))),
lookahead_tokens=int(os.getenv("DEEPINFER_USE_NEXTN_DECODE", "5")),
)
return llm

View File

@ -53,7 +53,7 @@ def main(dp_size: int, dp_rank: int) -> None:
max_model_len=32768,
max_num_seqs=4,
gpu_memory_utilization=0.85,
lookahead_tokens=min(1, int(os.getenv("DEEPINFER_USE_NEXTN_DECODE", "5"))),
lookahead_tokens=int(os.getenv("DEEPINFER_USE_NEXTN_DECODE", "5")),
)
from _prompts import OFFLINE_INPUT_LEN, OFFLINE_OUTPUT_LEN, lengthen_offline_prompts, prepare_offline_prompts