RFC: Agent-Ready Torch-TensorRT with a Local MCP Server

[zewenli98]

Summary

As an increasing number of people start vibe coding and as agents/tooling start driving optimization workflows, we need a stable, machine-readable compilation surface and a safe out-of-process execution path so optimization, validation, and benchmarking can be automated reliably and without fragile log-parsing or accidental unsafe execution.

This RFC proposes making Torch-TensorRT an agent-ready optimization system that agents can understand, inspect before compilation, execute, verify, and explain through stable, structured interfaces.

The implementation has two layers:

1. Extend Torch-TensorRT with machine-readable analysis, compilation reports, diagnostics, validation, benchmarking, and artifact manifests.
2. Expose these capabilities through an independent MCP Server running locally on the developer’s machine.

The initial scope is exclusively local. Remote servers, network APIs, distributed workers, multi-user authorization, and remote artifact storage are out of scope of this RFC, but could be considered in the future.

Motivation

Compiler information often reaches callers through stdout, stderr, or Python exceptions. Agents need a more reliable interface to determine:

• Whether a compilation request is valid.
• Which operators or partitions can be compiled.
• Where and why compilation failed.
• Whether a compiled model passes correctness checks.
• Whether it improves performance under a defined benchmark protocol.
• Which configuration and environment produced an artifact.

Structured interfaces should support both local agents and conventional Python tooling without requiring log parsing or MCP-specific compiler logic.

Goals

• Provide stable, versioned compilation reports and error categories.
• Enable preflight analysis before expensive engine building.
• Expose structured events and diagnostic artifacts.
• Provide official correctness-validation and benchmark interfaces.
• Record artifact identity, provenance, and compatibility constraints.
• Allow local agents to invoke these capabilities with explicit execution boundaries.

Non-Goals

• Remote MCP hosting or network-accessible services.
• Distributed compilation, remote GPU workers, or cloud orchestration.
• Multi-tenant authentication, authorization, or resource scheduling.
• Remote model upload or artifact storage.
• Making MCP a dependency of compilation or inference.
• Unrestricted agent-triggered Python or shell execution.
• Guaranteeing compilation success from preflight analysis.
• Treating validation on sampled inputs as universal correctness.

Architecture

Local AI Client / Coding Agent
              |
              | Local MCP connection
              v
Local Torch-TensorRT MCP Server
  - Tool exposure and input validation
  - Workspace and execution policy
  - Local job tracking
  - Report and artifact access
              |
              | Structured Python interfaces
              v
Torch-TensorRT
  - Preflight analysis
  - Compilation and diagnostics
  - Validation and benchmarking
  - Artifact metadata
              |
              v
Local Python environment, filesystem, and GPU

The proposed initial transport is stdio: the local MCP client launches the server as a subprocess. No network listener is required.

The server runs in an explicitly configured Python environment and uses the local GPU and filesystem. It remains a thin adapter over Torch-TensorRT capabilities rather than a separate compilation implementation.

1. Structured Compiler Interfaces

1.1 Versioned Compilation Reports

Each compilation invocation should produce a structured report, including unsuccessful invocations.

Required information:

Category	Contents
Identity	Schema version, job ID, status
Compiler	Torch-TensorRT version, frontend, backend
Environment	GPU, compute capability, driver, CUDA, TensorRT, PyTorch
Model	Graph hash, input signatures, shape profiles
Compilation	Duration, engine size, coverage, compiled and fallback segments
Diagnostics	Unsupported operators, warnings, categorized failures
Artifacts	Type, hash, size, local references

Coverage must have an explicit definition, including its denominator and excluded nodes. It must not be presented as a performance estimate.

1.2 Structured Failures

Failures should contain:

{
  "schema_version": "1.0",
  "status": "failed",
  "failure": {
    "category": "unsupported_operator",
    "stage": "lowering",
    "operator": "aten.some_op.default",
    "node_id": "node_47",
    "message": "No Dynamo converter is registered for this operator.",
    "retryable": false,
    "suggested_actions": [
      "Enable PyTorch fallback for this operator",
      "Provide a custom converter",
      "Rewrite the graph to avoid this operator"
    ]
  }
}

Initial taxonomy:

unsupported_operator
dynamic_shape_violation
invalid_input_profile
dtype_or_precision_conflict
graph_capture_failure
converter_failure
builder_failure
out_of_memory
runtime_load_failure
validation_failure
benchmark_failure
environment_incompatible
timeout
cancelled
internal_error

Retry guidance should distinguish retrying an unchanged request from retrying after modifying configuration or the environment.

1.3 Preflight Analysis

Provide an analyze() interface for inspecting a request before engine building.

It should report:

• Input specification and shape-profile issues.
• Detected unsupported operators.
• Expected compilation and fallback partitions, where determinable.
• Precision and environment constraints.
• Suggested actions.
• Checks completed, checks skipped, and unresolved conditions.

Analysis must distinguish confirmed findings from predictions. It should also disclose whether a check requires model execution rather than purely static inspection.

1.4 Events and Diagnostic Packages

Provide structured events alongside human-readable logs:

def on_event(event: dict) -> None:
    ...

Events should identify the job, stage, timestamp, and ordering. Examples include:

graph_captured
partition_completed
engine_build_started
engine_build_progress
engine_build_completed
artifact_saved
job_failed
job_cancelled

Emit progress percentages only when meaningful underlying measurements exist.

Persist diagnostic packages locally with structured reports, raw logs, and relevant graph artifacts. The MCP Server should expose current status and bounded diagnostic summaries without requiring agents to read entire logs.

2. Correctness Validation

Provide a validation harness with explicit input generation, metrics, and tolerances.

Illustrative API:

result = torch_tensorrt.validate(
    baseline=model,
    compiled=trt_model,
    input_generator=inputs,
    metrics=[
        "max_abs_error",
        "max_rel_error",
        "mean_rel_error",
    ],
    tolerances={
        "max_rel_error": 1e-3,
    },
)

Reports should include:

• Pass or fail status.
• Sample count and random seed, where applicable.
• Tested shapes and profile coverage.
• Numerical metrics and tolerance rules.
• NaN and infinity detection.
• Failed-case references.

The protocol must define output matching and relative-error behavior near zero. Passing validation demonstrates agreement for the tested inputs only.

3. Reproducible Benchmarking

Provide a benchmark harness that compares baseline and compiled execution under a declared protocol.

Illustrative API:

result = torch_tensorrt.benchmark(
    baseline=model,
    compiled=trt_model,
    input_spec=input_spec,
    warmup_iterations=100,
    iterations=5000,
    percentiles=[50, 95, 99],
    timing="cuda_events",
    synchronize=True,
)

Reports should record:

• Protocol version and timing semantics.
• Local hardware and software environment.
• Input shapes, precision, and execution configuration.
• Warmup and measurement iterations.
• Whether transfers and preprocessing are included.
• Baseline and compiled latency percentiles.
• Throughput measurement method.
• Speedup calculations.

Device execution time and end-to-end latency must remain distinct. Reports should record known local measurement conditions and avoid implying that the GPU was isolated unless isolation was established.

4. Local Artifact Manifests

Store artifacts and supporting evidence in a local output directory.

Illustrative layout:

job-output/
  manifest.json
  compile-report.json
  validation-report.json
  benchmark-report.json
  diagnostics/
  artifacts/

Artifact names and formats depend on compilation mode. Validation and benchmark reports are included only when those stages run; their absence does not imply success.

The manifest should record:

• Model, graph, and weights hashes.
• Compiler and environment versions.
• GPU architecture and compute capability.
• Frontend and compilation mode.
• Input profiles and precision.
• Fallback and plugin information.
• Artifact hashes and sizes.
• Compatibility constraints.
• Compilation timestamp and duration.
• Canonicalized compilation parameters.
• Relative references to supporting reports.

The local MCP Server should use artifact IDs associated with these manifests rather than requiring agents to infer artifact identity from filenames.

5. Local MCP Tool Surface

The following names are proposed interfaces, not claims about existing APIs.

Tool	Purpose
collect_environment	Inspect the configured local execution environment
analyze	Perform preflight analysis
compile	Start a local compilation job
get_job_status	Retrieve job state and progress
cancel_job	Request termination of a local job
validate	Compare a compiled artifact with its baseline
benchmark	Measure baseline and compiled performance
get_report	Retrieve a structured report
get_artifact_manifest	Inspect local artifact metadata

Long-running operations should return a job ID. Status queries should distinguish queued, running, succeeded, failed, cancellation-requested, and cancelled states.

Use a local worker subprocess for execution where appropriate. Define termination and cleanup behavior explicitly; requesting cancellation must not be reported as completed cancellation until the worker has stopped.

6. Local Execution Boundaries

Local-only operation does not mean unrestricted execution. The initial server should enforce:

• Explicitly configured workspace and output roots.
• Validated paths, including handling of symlinks and path traversal.
• Controlled model references or approved entry points.
• No generic shell or arbitrary-code execution tool.
• Timeouts and configurable job-concurrency limits.
• Explicit policies for overwriting and deleting files.
• Local request and outcome records.
• Bounded report responses, with full diagnostics retained as local files.

Model loading and model execution should be treated as execution-sensitive operations. The initial implementation should accept only developer-approved local models or entry points, rather than treating every file inside the workspace as trusted.

Remote credentials, network authorization, and remote artifact-transfer policies are outside this RFC.

7. Repository Location and Packaging

Decision: The local MCP adapter will live in the main Torch-TensorRT repository but will be distributed as a separate, optional Python package.

Repository location and packaging are independent choices: the repository determines how development is coordinated, while packaging determines what users install. Co-location allows a single change to update core contracts, the adapter, tests, and documentation while those contracts are still evolving. Separate packaging keeps MCP dependencies out of the core compiler installation.

Component	Location	Distribution / Responsibility
Structured compiler APIs and report schemas	Core source tree	torch-tensorrt; compiler-owned contracts
Validation, benchmarking, and manifests	Core or supporting tooling in the repository	Usable without MCP
Local MCP adapter	Separate directory in the same repository	Optional torch-tensorrt-mcp package (proposed name)
MCP dependencies and entry point	Adapter package	Not required by the core package
Adapter integration tests	Same repository	Verify compatibility with core interfaces

Proposed layout:

Torch-TensorRT/
├── py/
│   └── torch_tensorrt/        # Compiler APIs and contracts
├── tools/
│   └── mcp/
│       ├── pyproject.toml     # Separate package metadata
│       ├── src/
│       │   └── torch_tensorrt_mcp/
│       └── tests/
└── docs/                      # Local MCP setup and usage

Boundaries

• One-way dependency: torch-tensorrt-mcp depends on public Torch-TensorRT APIs; torch-tensorrt never depends on MCP. The adapter translates MCP requests into core API calls and returns their structured results. It must not duplicate operator analysis, rely on compiler internals, or reconstruct reports by parsing logs.
• Independent installation, coordinated compatibility: The adapter is published separately with an explicitly tested Torch-TensorRT compatibility range. Releases are coordinated with core API changes initially; adapter-only releases are allowed once contracts stabilize. Version numbers need not match; a compatibility table is published instead.
• No silent environment changes: The server runs from an explicitly selected Python environment. Installing or launching the adapter must not upgrade Torch-TensorRT, PyTorch, or TensorRT; incompatible environments produce an actionable error.
• Layered testing:
  • Core contract tests: reports, schemas, taxonomy, and API behavior.
  • Adapter tests without GPU execution: request validation, serialization, workspace boundaries, and job state, using test doubles.
  • GPU integration tests: real analysis, compilation, validation, and benchmarking.

Future Repository Split

A separate repository may be considered once the adapter has independent maintainers, a different release cadence, or support for multiple optimization backends. Keeping a clean package boundary now preserves that option.

8. User Workflow Example

This section illustrates the intended end-to-end experience after the capabilities in this RFC are implemented. All commands, tool names, and outputs are proposed designs, not existing interfaces.

8.1 One-Time Setup

Install the optional adapter into an existing Torch-TensorRT environment:

conda activate trt-env          # Already contains torch, torch-tensorrt, and TensorRT
pip install torch-tensorrt-mcp  # Optional package; does not upgrade compiler dependencies

Register the local server with an MCP client using stdio transport (no network listener):

{
  "mcpServers": {
    "torch-tensorrt": {
      "command": "/home/alice/miniconda3/envs/trt-env/bin/torch-tensorrt-mcp",
      "args": [
        "--workspace", "/home/alice/projects/detector",
        "--output-root", "/home/alice/projects/detector/trt-jobs",
        "--allowed-entrypoint", "models/build.py:build_model"
      ]
    }
  }
}

The configuration explicitly defines the Python environment, the readable and writable directories, and the approved model entry point.

8.2 Example: Optimizing a Detection Model with FP16

The user asks the coding agent:

Compile the detection model returned by build_model to TensorRT with FP16, input shape [1,3,640,640]. Keep the maximum relative error within 1e-3, and tell me the speedup on this L4.

The agent then performs the following steps.

Step 1: Inspect the environment. collect_environment returns the GPU (L4, compute capability 8.9) and the CUDA, TensorRT, and PyTorch versions. An incompatible environment returns environment_incompatible before any compilation work starts.

Step 2: Run preflight analysis. analyze(entrypoint, input_spec, precision="fp16") returns:

{
  "status": "warnings",
  "unsupported_ops": [
    {"operator": "aten.nonzero.default", "node_id": "node_312",
     "confirmed": true, "reason": "data-dependent output shape"}
  ],
  "predicted_partitions": {"trt": 2, "fallback": 1},
  "suggested_actions": ["Allow PyTorch fallback for aten.nonzero.default"]
}

The agent explains that nonzero in post-processing will fall back to PyTorch, producing two TensorRT segments and one PyTorch segment, and asks the user whether this fallback is acceptable.

Step 3: Start compilation as a local job. After user approval, compile(...) immediately returns a job_id. The agent polls get_job_status and observes structured events:

graph_captured → partition_completed (trt_nodes=417, fallback_nodes=11)
→ engine_build_started (segment 0) → engine_build_completed (segment 0) → ...

The agent can report precisely that it is building the second TensorRT engine. The user may request cancel_job at any time.

Step 4: Validate correctness. validate(job_id, tolerances={"max_rel_error": 1e-3}) returns:

{"status": "pass", "samples": 200,
 "metrics": {"max_rel_error": 0.00064}, "nan_or_inf_detected": false}

The agent states that the result demonstrates agreement on the 200 tested inputs only.

Step 5: Benchmark. benchmark(job_id, iterations=5000, timing="cuda_events") reports p95 latency decreasing from 23.9 ms to 9.6 ms (about 2.5x p95 speedup), excluding transfers and preprocessing.

Step 6: Deliver the artifact. After get_artifact_manifest(job_id), the agent summarizes:

Compilation succeeded with FP16 and one PyTorch fallback segment (nonzero). Validation passed (max relative error 6.4e-4). p95 speedup is about 2.5x. The artifact and all reports are in trt-jobs/<job_id>/; the manifest records model hashes, environment versions, input profiles, and compatibility constraints.

8.3 Failure Handling

If engine building runs out of memory, get_job_status returns:

{"status": "failed",
 "failure": {"category": "out_of_memory", "stage": "engine_build",
             "retryable": true,
             "suggested_actions": ["Reduce workspace size", "Reduce max batch in shape profile"]}}

Based on the category, the agent determines that an unchanged retry is not useful, adjusts the configuration (for example, reducing workspace size), recompiles, and reports the change to the user without reading raw logs.

8.4 Resulting Change in User Experience

Before	After
Write compilation scripts, read logs, and benchmark manually	State goals and constraints in natural language
Search stderr to diagnose failures	Receive categorized failures and suggested actions
Speedup depends on ad hoc measurement	Speedup follows a declared, traceable benchmark protocol
Identify artifacts by filename	Identify artifacts by manifest, provenance, and compatibility

The entire workflow runs locally. The same analyze, validate, and benchmark interfaces remain usable directly from Python scripts and CI without MCP.

Delivery Plan

Phase 1: Structured Core Capabilities

Implement compilation reports, error taxonomy, preflight analysis, events, and diagnostic packages.

Phase 2: Verification and Artifact Evidence

Implement validation, benchmarking, reproducibility metadata, and local manifests.

Phase 3: Local MCP Adapter

Implement a stdio-based server as the optional torch-tensorrt-mcp package in the main repository, starting with environment inspection, analysis, and report retrieval. Add local compilation, validation, benchmarking, and cancellation after the execution contract is defined.

Success Criteria

• Agents classify failures without parsing free-form logs.
• Preflight predictions remain distinguishable from observed results.
• Long-running jobs expose reliable state and diagnostic information.
• Validation and benchmark results identify their test conditions.
• Artifacts can be traced to their inputs, configuration, and environment.
• Local execution stays within configured filesystem and resource boundaries.
• The same structured interfaces work from Python tooling without MCP.
• Installing torch-tensorrt never pulls in MCP dependencies.

[zewenli98]

With this feature implemented, we can even ask agents to do more like:

1. write converters or lowering passes on the fly (we've already exposed API for users to add converters and lowering passes)
2. write TRT plugins
3. generate kernels w/o or w/ KerGen tools
4. ...

[narendasan]

Lets use dataloaders and other pytorch infra to make this a bit simpler to use.

[bowang007]

The overall proposal of the pipeline looks good and makes sense from users' perspective.
I would suggest:

1. Refine this proposal if you plan to have a very detailed roadmap, we could have a google doc so that everyone in this team could leave the comments and you could add more details for more rounds. Or you could just propose a MVP with more details for now. The current status of this proposal is general but contains some unnecessary details like the repo split plan for the long term.
2. A lot of tools are mentioned in this proposal, it would be perfect if we could mention the functionalities of each one and try to merge some of them to make it simpler and concise for maintenance as well.
3. Besides all of the capabilities that are mentioned in this proposal, what makes this agent different from or distinguished from other general agent like claude code or codex nowadays? Some of tools and the capabilities of the proposal agent could be achieved by claude code/codex, such as collect_environment, preflight analysis, compile, these are several lines of command jobs that could be conducted by claude code/codex as well. I would propose you to do something more like:

• Collect the environment, check any dependent library is missing, upgrade TRT and so on to make it standout, and a bunch of details need to be provided with the tools to make it stand out of the general agents job.
• Preflight analysis, analyze the model and environment then adjust the dependent library, and then propose many different compilation strategies to try it out.
• Compile, in the compile tool, maybe we could introduce more knobs in this compile progress, like compile multiple different backends for benchmarking and so on.

In the long term, we could target it as a Torch-TensorRT Compiler Specialist Agent or TRT compiler Specialist Agent that could take many different front ends of models and compile it to SOL perf on various platform and beat the other agents who only does some normal optimization.
cc: @narendasan