diff --git a/launcher/src/main.rs b/launcher/src/main.rs index 26a2e4aa9..748680a6a 100644 --- a/launcher/src/main.rs +++ b/launcher/src/main.rs @@ -350,6 +350,10 @@ struct Args { #[clap(long, env)] speculative_tokens: Option, + // The maximum batch size past which speculative decoding is disabled. + #[clap(default_value = "32", long, env)] + speculation_max_batch_size: usize, + /// The list of adapter ids to preload during initialization (to avoid cold start times). #[clap(long, env)] preloaded_adapter_ids: Vec, @@ -638,6 +642,7 @@ fn shard_manager( quantize: Option, compile: bool, speculative_tokens: Option, + speculation_max_batch_size: usize, preloaded_adapter_ids: Vec, preloaded_adapter_source: Option, predibase_api_token: Option, @@ -802,6 +807,12 @@ fn shard_manager( envs.push(("CHUNKED_PREFILL".into(), chunked_prefill.to_string().into())); } + // Speculative decoding max batch size + envs.push(( + "LORAX_SPECULATION_MAX_BATCH_SIZE".into(), + speculation_max_batch_size.to_string().into(), + )); + // Backend if backend == Backend::FlashInfer { envs.push(("FLASH_INFER".into(), "1".into())); @@ -1244,6 +1255,7 @@ fn spawn_shards( let quantize = args.quantize; let compile = args.compile; let speculative_tokens = args.speculative_tokens; + let speculation_max_batch_size = args.speculation_max_batch_size; let preloaded_adapter_ids = args.preloaded_adapter_ids.clone(); let preloaded_adapter_source = args.preloaded_adapter_source.clone(); let predibase_api_token = args.predibase_api_token.clone(); @@ -1271,6 +1283,7 @@ fn spawn_shards( quantize, compile, speculative_tokens, + speculation_max_batch_size, preloaded_adapter_ids, preloaded_adapter_source, predibase_api_token,