feat(nemotron-h): add native Edge execution

families/nemotron_h/build_request.py

@@ -0,0 +1,98 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""nemotron_h 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:
+    """nemotron_h-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: ClassVar[tuple[int, ...]] = ()
+    dynamic_kv_cache: ClassVar[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,
+        "fp32_layers": (),
+        "dynamic_kv_cache": False,
+    }
+    for name, default in unsupported.items():
+        value = getattr(request, name, default)
+        if name == "quantization" and value == "none":
+            continue
+        if name == "fp32_layers" and isinstance(value, (list, tuple)) and not value:
+            continue
+        if value != default:
+            raise NotImplementedError(f"nemotron_h does not support {name}")
+    names = {field.name for field in fields(BuildRequest)}
+    if unknown := set(vars(request)) - names - set(unsupported):
+        raise ValueError(f"unknown nemotron_h build inputs: {sorted(unknown)}")
+    return BuildRequest(**{name: getattr(request, name) for name in names})

families/nemotron_h/cli.json

@@ -0,0 +1,120 @@
+{
+  "version": 1,
+  "commands": [
+    {
+      "name": "build",
+      "help": "Build one nemotron_h 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"
+          ],
+          "default": "fp32"
+        },
+        {
+          "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": "execution_variant",
+          "flags": [
+            "--execution-variant"
+          ],
+          "type": "string",
+          "choices": [
+            "dflash"
+          ],
+          "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"
+        }
+      ]
+    }
+  ]
+}

families/nemotron_h/cli.py

@@ -0,0 +1,59 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""nemotron_h-owned build command and typed inputs; importing this module is CPU-only."""
+
+from __future__ import annotations
+
+from dataclasses import replace
+from pathlib import Path
+
+from tensorrt_model_connect.build import select_backend
+from tensorrt_model_connect.bundle_writer import BundleWriter
+from tensorrt_model_connect.graph_transform import graph_transform
+from tensorrt_model_connect.model_support import load_model_metadata, resolve_family, resolve_model
+
+from .build_request import BuildRequest
+
+from .edge_llm.cli import execution_inputs
+from .edge_llm.config import with_execution
+
+
+def build_bundle(request: BuildRequest, output: Path) -> None:
+    """Build and publish through the owning model, preserving atomic failure."""
+    request = replace(request, output_path=output)
+    select_backend(request.backend)
+    from .model import build as build_model
+
+    writer = BundleWriter(output)
+    try:
+        with graph_transform(request.graph_transform):
+            build_model(request, writer)
+        writer.finish()
+    except BaseException:
+        writer.abort()
+        raise
+
+def build(
+    *, model: str, output: Path, revision: str | None = None,
+    task: str = "text_generation", precision: str = "fp32", backend: str = "trt",
+    max_sequence_length: int | None = None, tensor_parallel_size: int = 1,
+    verbose: bool = False,
+    execution_variant: str | None = None, companion: list[str] | tuple[str, ...] = (),
+) -> int:
+    """Run the declared owner command; help never imports this handler."""
+    if precision not in {"fp32", "fp16"}:
+        raise ValueError("Nemotron-H precision must be fp32 or fp16")
+    execution = execution_inputs(execution_variant, companion)
+    model_dir = resolve_model(model, revision)
+    resolve_family(load_model_metadata(model_dir), "nemotron_h")
+    request = BuildRequest(
+        model_dir=model_dir, output_path=output, family="nemotron_h",
+        task=task, precision=precision, backend=backend,
+        max_sequence_length=max_sequence_length, tensor_parallel_size=tensor_parallel_size,
+        verbose=verbose,
+    )
+    if execution is not None:
+        request = with_execution(request, execution)
+    build_bundle(request, output)
+    return 0