feat(llama): add native Edge execution

apps/benchmark/trtmc_benchmark/tests/test_benchmark.py

@@ -1032,7 +1032,9 @@ def test_build_command_passes_manifest_backend_and_dynamic_kv_cache(tmp_path: Pa
 
     command = _build_command(model, tmp_path / "checkpoint", tmp_path / "model.bundle", (case,))
 
-    assert command[command.index("--backend") + 1] == "trt_rtx"
+    # Declared family commands serialize scalar flags as --name=value.
+    if "--backend=trt_rtx" not in command:
+        assert command[command.index("--backend") + 1] == "trt_rtx"
     assert command.count("--dynamic-kv-cache") == 1
 
 

families/llama/build_request.py

@@ -0,0 +1,94 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""llama build inputs and strict compatibility for existing Python callers."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, fields
+from pathlib import Path
+import re
+from typing import ClassVar
+
+from tensorrt_model_connect.graph_transform import GraphTransform
+
+
+def _validate_id(field: str, value: str) -> None:
+    if not isinstance(value, str) or re.fullmatch(r"[a-z][a-z0-9_]*", value) is None:
+        raise ValueError(f"{field} must be a lowercase identifier")
+
+
+@dataclass(frozen=True)
+class BuildRequest:
+    """llama-owned inputs; unsupported legacy controls are read-only defaults."""
+
+    model_dir: Path
+    output_path: Path
+    family: str
+    task: str
+    precision: str
+    backend: str = "trt"
+    max_sequence_length: int | None = None
+    image_height: ClassVar[int | None] = None
+    image_width: ClassVar[int | None] = None
+    video_num_frames: ClassVar[int | None] = None
+    max_batch_size: ClassVar[int] = 1
+    tensor_parallel_size: int = 1
+    context_parallel_size: ClassVar[int] = 1
+    quantization: ClassVar[str | None] = None
+    fp32_layers: tuple[int, ...] = ()
+    dynamic_kv_cache: bool = False
+    verbose: bool = False
+    graph_transform: GraphTransform | None = None
+
+    def __post_init__(self) -> None:
+        if not self.precision:
+            raise ValueError("precision must be non-empty")
+        _validate_id("family", self.family)
+        _validate_id("task", self.task)
+        if self.backend not in {"trt", "trt_rtx"}:
+            raise ValueError("backend must be 'trt' or 'trt_rtx'")
+        if self.max_sequence_length is not None and self.max_sequence_length < 1:
+            raise ValueError("max_sequence_length must be positive")
+        for field in ("image_height", "image_width", "video_num_frames"):
+            value = getattr(self, field)
+            if value is not None and value < 1:
+                raise ValueError(f"{field} must be positive")
+        if self.max_batch_size < 1:
+            raise ValueError("max_batch_size must be positive")
+        if self.tensor_parallel_size < 1:
+            raise ValueError("tensor_parallel_size must be positive")
+        if self.context_parallel_size < 1:
+            raise ValueError("context_parallel_size must be positive")
+        if self.quantization is not None and not self.quantization:
+            raise ValueError("quantization must be non-empty when provided")
+        if any(layer < 0 for layer in self.fp32_layers):
+            raise ValueError("fp32_layers must contain non-negative indices")
+        if not isinstance(self.dynamic_kv_cache, bool):
+            raise ValueError("dynamic_kv_cache must be a bool")
+        if self.graph_transform is not None and not callable(self.graph_transform):
+            raise ValueError("graph_transform must be callable when provided")
+
+
+def coerce_request(request: object) -> BuildRequest:
+    """Reject unsupported/unknown legacy inputs before converting to owner fields."""
+    if isinstance(request, BuildRequest):
+        return request
+    unsupported = {
+        "image_height": None,
+        "image_width": None,
+        "video_num_frames": None,
+        "max_batch_size": 1,
+        "context_parallel_size": 1,
+        "quantization": None,
+    }
+    for name, default in unsupported.items():
+        value = getattr(request, name, default)
+        if name == "quantization" and value == "none":
+            continue
+        if value != default:
+            raise NotImplementedError(f"llama does not support {name}")
+    names = {field.name for field in fields(BuildRequest)}
+    if unknown := set(vars(request)) - names - set(unsupported):
+        raise ValueError(f"unknown llama build inputs: {sorted(unknown)}")
+    return BuildRequest(**{name: getattr(request, name) for name in names})

families/llama/cli.json

@@ -1,21 +1,235 @@
 {
   "version": 1,
-  "commands": [{
-    "name": "build-speculative",
-    "help": "Build the FP16 Llama 3.1 / EAGLE3 prototype bundle",
-    "executor": "python",
-    "handler": "speculative.build:build_speculative",
-    "arguments": [
-      {"name": "model_dir", "flags": ["--model-dir"], "type": "path", "required": true},
-      {"name": "draft_dir", "flags": ["--draft-dir"], "type": "path", "required": true},
-      {"name": "output", "flags": ["--output"], "type": "path", "required": true},
-      {"name": "spec_dec", "flags": ["--spec-dec"], "type": "string", "choices": ["eagle3"], "default": "eagle3"},
-      {"name": "max_sequence_length", "flags": ["--max-sequence-length"], "type": "int", "default": 2048},
-      {"name": "max_query", "flags": ["--max-query"], "type": "int", "default": 64},
-      {"name": "execution_profiles", "flags": ["--execution-profiles"], "type": "string", "choices": ["single", "split"], "default": "single"},
-      {"name": "prefill_query", "flags": ["--prefill-query"], "type": "int", "default": 64},
-      {"name": "draft_depth", "flags": ["--draft-depth"], "type": "int", "default": 4},
-      {"name": "verbose", "flags": ["--verbose"], "type": "bool", "action": "store_true", "default": false}
-    ]
-  }]
+  "commands": [
+    {
+      "name": "build",
+      "help": "Build one llama TensorRT bundle",
+      "executor": "python",
+      "handler": "cli:build",
+      "arguments": [
+        {
+          "name": "model",
+          "type": "string",
+          "help": "Hugging Face model ID or local snapshot"
+        },
+        {
+          "name": "output",
+          "flags": [
+            "-o",
+            "--output"
+          ],
+          "type": "path",
+          "required": true
+        },
+        {
+          "name": "revision",
+          "flags": [
+            "--revision"
+          ],
+          "type": "string"
+        },
+        {
+          "name": "task",
+          "flags": [
+            "--task"
+          ],
+          "type": "string",
+          "choices": [
+            "text_generation"
+          ],
+          "default": "text_generation"
+        },
+        {
+          "name": "precision",
+          "flags": [
+            "--precision"
+          ],
+          "type": "string",
+          "choices": [
+            "fp32",
+            "fp16",
+            "bf16"
+          ],
+          "help": "Compute precision (default: fp16 for paired execution, fp32 otherwise)"
+        },
+        {
+          "name": "backend",
+          "flags": [
+            "--backend"
+          ],
+          "type": "string",
+          "choices": [
+            "trt",
+            "trt_rtx"
+          ],
+          "default": "trt"
+        },
+        {
+          "name": "max_sequence_length",
+          "flags": [
+            "--max-sequence-length"
+          ],
+          "type": "int"
+        },
+        {
+          "name": "tensor_parallel_size",
+          "flags": [
+            "--tensor-parallel-size"
+          ],
+          "type": "int",
+          "choices": [
+            1,
+            2,
+            4,
+            8
+          ],
+          "default": 1
+        },
+        {
+          "name": "verbose",
+          "flags": [
+            "--verbose"
+          ],
+          "type": "bool",
+          "action": "store_true",
+          "default": false
+        },
+        {
+          "name": "fp32_layers",
+          "flags": [
+            "--fp32-layer"
+          ],
+          "type": "int",
+          "action": "append",
+          "default": []
+        },
+        {
+          "name": "dynamic_kv_cache",
+          "flags": [
+            "--dynamic-kv-cache"
+          ],
+          "type": "bool",
+          "action": "store_true",
+          "default": false
+        },
+        {
+          "name": "execution_variant",
+          "flags": [
+            "--execution-variant"
+          ],
+          "type": "string",
+          "choices": [
+            "eagle3"
+          ],
+          "help": "Explicit paired execution mode; no automatic draft discovery"
+        },
+        {
+          "name": "companion",
+          "flags": [
+            "--companion"
+          ],
+          "type": "string",
+          "action": "append",
+          "default": [],
+          "help": "Local companion checkpoint as ROLE=LOCAL_DIR"
+        }
+      ]
+    },
+    {
+      "name": "build-speculative",
+      "help": "Build the FP16 Llama 3.1 / EAGLE3 prototype bundle",
+      "executor": "python",
+      "handler": "speculative.build:build_speculative",
+      "arguments": [
+        {
+          "name": "model_dir",
+          "flags": [
+            "--model-dir"
+          ],
+          "type": "path",
+          "required": true
+        },
+        {
+          "name": "draft_dir",
+          "flags": [
+            "--draft-dir"
+          ],
+          "type": "path",
+          "required": true
+        },
+        {
+          "name": "output",
+          "flags": [
+            "--output"
+          ],
+          "type": "path",
+          "required": true
+        },
+        {
+          "name": "spec_dec",
+          "flags": [
+            "--spec-dec"
+          ],
+          "type": "string",
+          "choices": [
+            "eagle3"
+          ],
+          "default": "eagle3"
+        },
+        {
+          "name": "max_sequence_length",
+          "flags": [
+            "--max-sequence-length"
+          ],
+          "type": "int",
+          "default": 2048
+        },
+        {
+          "name": "max_query",
+          "flags": [
+            "--max-query"
+          ],
+          "type": "int",
+          "default": 64
+        },
+        {
+          "name": "execution_profiles",
+          "flags": [
+            "--execution-profiles"
+          ],
+          "type": "string",
+          "choices": [
+            "single",
+            "split"
+          ],
+          "default": "single"
+        },
+        {
+          "name": "prefill_query",
+          "flags": [
+            "--prefill-query"
+          ],
+          "type": "int",
+          "default": 64
+        },
+        {
+          "name": "draft_depth",
+          "flags": [
+            "--draft-depth"
+          ],
+          "type": "int",
+          "default": 4
+        },
+        {
+          "name": "verbose",
+          "flags": [
+            "--verbose"
+          ],
+          "type": "bool",
+          "action": "store_true",
+          "default": false
+        }
+      ]
+    }
+  ]
 }