A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 09:45 UTC
Update default neuron config for speculation (#18274)
Signed-off-by: Elaine Zhao <elaineyz@amazon.com> Co-authored-by: Shashwat Srijan <sssrijan@amazon.com> Co-authored-by: Aakash Shetty <sheaak@amazon.com>
aws-elaineyz committed
May 22, 2025 at 02:18 UTC
ebed81fbf5e4549b26b6baf74cea0fe4551dd915
1 file changed
+3
-1
vllm/model_executor/model_loader/neuronx_distributed.py
+3
-1
@@ -502,7 +502,7 @@ def _get_default_neuron_config(model_config: ModelConfig,
502
max_context_length=scheduler_config.max_model_len,
503
seq_len=scheduler_config.max_model_len,
504
enable_bucketing=True,
505
- is_continuous_batching=(batch_size > 1),
505
+ is_continuous_batching=True,
506
quantized=False,
507
torch_dtype=TORCH_DTYPE_TO_NEURON_AMP[model_config.dtype],
508
padding_side="right",
@@ -520,6 +520,7 @@ def _get_default_speculation_config(model_config: ModelConfig,
520
args."""
521
neuron_config = dict(
522
tp_degree=parallel_config.tensor_parallel_size,
523
+ ctx_batch_size=1,
524
batch_size=scheduler_config.max_num_seqs,
525
max_context_length=scheduler_config.max_model_len,
526
seq_len=scheduler_config.max_model_len,
@@ -527,6 +528,7 @@ def _get_default_speculation_config(model_config: ModelConfig,
528
trace_tokengen_model=False,
529
enable_fused_speculation=True,
530
enable_bucketing=True,
531
+ is_continuous_batching=True,
532
quantized=False,
533
torch_dtype=TORCH_DTYPE_TO_NEURON_AMP[model_config.dtype],
534
on_device_sampling_config=dict(