diff --git a/launcher/src/main.rs b/launcher/src/main.rs index 26a2e4aa9..df707c00f 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(long, env)] + speculation_max_batch_size: Option, + /// 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: Option, preloaded_adapter_ids: Vec, preloaded_adapter_source: Option, predibase_api_token: Option, @@ -802,6 +807,14 @@ fn shard_manager( envs.push(("CHUNKED_PREFILL".into(), chunked_prefill.to_string().into())); } + // Speculative decoding max batch size + if let Some(speculation_max_batch_size) = speculation_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 +1257,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 +1285,7 @@ fn spawn_shards( quantize, compile, speculative_tokens, + speculation_max_batch_size, preloaded_adapter_ids, preloaded_adapter_source, predibase_api_token,