Add TRTLLM compilation + speculative decoding validation.

Author: Joseph Zhang

src/sagemaker/serve/builder/model_builder.py

@@ -1283,7 +1283,7 @@ def _model_builder_optimize_wrapper(
             # TRTLLM is used by Neo if the following are provided:
             #  1) a GPU instance type
             #  2) compilation config
-            gpu_instance_families = ["g5", "g6", "p4d", "p5"]
+            gpu_instance_families = ["g5", "g6", "p4d", "p4de", "p5"]
             is_gpu_instance = optimization_instance_type and any(
                 gpu_instance_family in optimization_instance_type
                 for gpu_instance_family in gpu_instance_families
@@ -1296,8 +1296,16 @@ def _model_builder_optimize_wrapper(
                 keyword in self.model.lower() for keyword in llama_3_1_keywords
             )
 
-            if is_gpu_instance and self.model and is_llama_3_1 and self.is_compiled:
-                raise ValueError("Compilation is not supported for Llama-3.1 with a GPU instance.")
+            if is_gpu_instance and self.model and self.is_compiled:
+                if is_llama_3_1:
+                    raise ValueError(
+                        "Compilation is not supported for Llama-3.1 with a GPU instance."
+                    )
+                if speculative_decoding_config:
+                    raise ValueError(
+                        "Compilation is not supported with speculative decoding with "
+                        "a GPU instance."
+                    )
 
             self.sagemaker_session.sagemaker_client.create_optimization_job(**input_args)
             job_status = self.sagemaker_session.wait_for_optimization_job(job_name)

tests/unit/sagemaker/serve/builder/test_model_builder.py

@@ -2891,3 +2891,58 @@ def test_optimize_with_gpu_instance_and_llama_3_1_and_compilation(
                 output_path="s3://bucket/code/",
             ),
         )
+
+    @patch.object(ModelBuilder, "_prepare_for_mode")
+    @patch.object(ModelBuilder, "_get_serve_setting", autospec=True)
+    def test_optimize_with_gpu_instance_and_compilation_with_speculative_decoding(
+        self,
+        mock_get_serve_setting,
+        mock_prepare_for_mode,
+    ):
+        mock_prepare_for_mode.side_effect = lambda *args, **kwargs: (
+            {
+                "S3DataSource": {
+                    "CompressionType": "None",
+                    "S3DataType": "S3Prefix",
+                    "S3Uri": "s3://bucket/code/code/",
+                }
+            },
+            {"DTYPE": "bfloat16"},
+        )
+
+        mock_pysdk_model = Mock()
+        mock_pysdk_model.model_data = None
+        mock_pysdk_model.env = {"HF_MODEL_ID": "modelid"}
+
+        sample_input = {"inputs": "dummy prompt", "parameters": {}}
+
+        sample_output = [{"generated_text": "dummy response"}]
+
+        dummy_schema_builder = SchemaBuilder(sample_input, sample_output)
+
+        model_builder = ModelBuilder(
+            model="modelid",
+            schema_builder=dummy_schema_builder,
+            env_vars={"HF_TOKEN": "token"},
+            model_metadata={
+                "CUSTOM_MODEL_PATH": "s3://bucket/path/",
+            },
+            role_arn="role-arn",
+            instance_type="ml.g5.2xlarge",
+        )
+
+        model_builder.pysdk_model = mock_pysdk_model
+
+        self.assertRaisesRegex(
+            ValueError,
+            "Compilation is not supported with speculative decoding with a GPU instance.",
+            lambda: model_builder.optimize(
+                job_name="job_name-123",
+                speculative_decoding_config={
+                    "ModelProvider": "custom",
+                    "ModelSource": "s3://data-source",
+                },
+                compilation_config={"OverrideEnvironment": {"OPTION_TENSOR_PARALLEL_DEGREE": "2"}},
+                output_path="s3://bucket/code/",
+            ),
+        )