[neo] Allow setting GPU memory-related LMI options in sharding jobs

serving/docker/partition/sm_neo_shard.py

@@ -95,14 +95,27 @@ def copy_non_safetensors_files(self, input_dir: str, output_dir: str):
     def shard_lmi_dist_model(self, input_dir: str, output_dir: str,
                              pp_degree: int, tp_degree: int,
                              chunk_mb: int) -> None:
+        # For engine args which can affect GPU memory utilization, use LMI defaults
+        # unless specified otherwise by the customer
+        gpu_memory_utilization = float(
+            self.properties.get("option.gpu_memory_utilization", 0.9))
+        enforce_eager: bool = str(
+            self.properties.get("option.enforce_eager",
+                                False)).lower() == "true"
+        max_rolling_batch_size = int(
+            self.properties.get("option.max_rolling_batch_size", 256))
+        max_model_len = int(self.properties.get("option.max_model_len", None))
 
         engine_args = VllmEngineArgs(
             model=input_dir,
             pipeline_parallel_size=pp_degree,
             tensor_parallel_size=tp_degree,
-            enforce_eager=True,
             disable_custom_all_reduce=True,
             distributed_executor_backend="mp",
+            gpu_memory_utilization=gpu_memory_utilization,
+            enforce_eager=enforce_eager,
+            max_num_seqs=max_rolling_batch_size,
+            max_model_len=max_model_len,
         )
         engine = engine_from_args(engine_args)