← All blogs

Engineering · 8 min read

How an AI agent adds new models to our ultra-fast inference pipeline

What changes when the work of porting a new VLA/WAM into a 30 ms robot stack goes from a 3–5 day engineering task to a 2–3 hour automated run with explicit correctness and latency gates.

Where the Model Integration Agent fits in the broader system

Adding a new model to a running system is a four-agent pipeline. The Model Integration Agent sits in the middle and produces the correctness + latency contract that downstream agents depend on.

INPUTModel & Requirements& ArtifactsCheckpoint (model weights)Reference inference codeExample dataset / oracleLatency targetHardware targets★ MAIN TOPICModel IntegrationAgentscout → spec → pipeline→ frontend → register→ validate → benchmark→ docs3 guardrails · 3 retriesalignment ≥ 0.99 · 0 NaN · ≤ targetPARALLEL · MANDATORYModel-specific KernelOptimization AgentHot-path tuning for this modelPARALLEL · OPTIONALGeneral KernelImplementation AgentFor unsupported ops (rare)FOLLOW-UPFull SR Recovery Agentv.s. the human demo SRthreshold. Recovers fromKV-cache bugs, cross-attninit, action expert issues.Optimized Runtime — ready to run, meets requirement
The full VLA deployment pipeline. The Model Integration Agent (highlighted) is the keystone — it must produce alignment ≥ 0.99, zero NaN, and meet the latency budget before downstream agents can run.

The Model Integration Agent (highlighted) is the keystone. Downstream agents cannot run until it produces an integration that meets alignment ≥ 0.99 against the reference PyTorch implementation, runs zero NaNs across a 20-shot soak, and stays under the requested end-to-end latency budget.

Why automated model integration is hard

Every model integration is a structured bookkeeping problem, not a kernel engineering problem. The shared low-level library exposes hundreds of compute primitives — matrix multiplications, normalizations, attention calls, fused activations. The per-model work is roughly 1,200 lines that translate a model's forward pass into those primitives, plus a 200-line wrapper for weight loading, buffer allocation, and GPU-graph capture.

The canonical pattern across every modern VLA looks the same:

  • Run the vision encoder — a small transformer that turns the camera image into a token list.
  • Project the vision tokens into the language model's space and concatenate with the text prompt.
  • Run the language model once over the prefix to build a key/value cache.
  • Run an action expert for 10 denoising steps that produces the next chunk of robot actions.
  • Capture the entire forward into a single GPU-graph replay so per-call overhead is essentially zero.

Four things make this hard to do by hand and equally hard to do without discipline:

  1. Silent failure modes. The most expensive bug in the manual SmolVLA integration was a wrong weight layout that produced all-NaN action outputs while the program ran end-to-end with no error.
  2. Shape-dependent pieces. Layers, hidden dimension, head ratio, FFN size, attention mask pattern, FP8 vs FP16 path — every source-model change chews through a different combination of patterns.
  3. GPU-graph capture is unforgiving. No memory copies, no Python branches, no buffer allocation after the first call.
  4. Kernel selection depends on shape and hardware. The shared library exposes several competing primitives for the same op. The right one for a given shape on a given card is not always obvious. In SmolVLA, an FP8 GEMM path that compiled fine produced alignment 0.003; switching to the cuBLASLt-backed path restored 0.998.

An LLM-based agent has the same familiar failure modes: hallucinated layer names, transposed masks, wrong buffer ownership, picked the wrong kernel. The agent must do the integration with the discipline a careful engineer would apply, and be auditable when it fails.

Output alignment ≥ 0.99 with the reference on the same inputs. 0 NaN across 20 inferences. Action chunk at the expected shape. End-to-end latency is a deliverable — if the budget is missed, retry one knob at a time, then escalate to a human engineer with a structured hand-off. If any of these fail, the agent writes a block report and stops. Broken code is not committed.

Inside the Model Integration Agent

The agent is an orchestrator with eight sub-agents — one per step in the canonical workflow — and a hard gate between each. The orchestrator carries an 8-step workflow plus a known-failure-modes note from prior integrations.

Model Integration Agent — orchestratorsystem prompt: 8-step workflow + known failure modes from past integrations1 · scoutscan_checkpointscan_reference_repo→ INTEG PLAN.mdplan exists2 · specwrite_spec_rtxwrite_spec_thor→ _spec.py ×2all keys load3 · pipelinewrite_pipeline_rtxwrite_pipeline_thor→ pipeline_*.py ×2finite on dummy4 · frontendwrite_frontend_rtxwrite_frontend_thor→ frontend_*.py ×2load_model OK5 · register_PIPELINE_MAPapi.py allow-list→ 5 file editsresolves6 · validaterun PyTorch oracle20-shot NaN soak · cos→ alignment★ alignment ≥ 0.997 · benchmarkmeasure e2e P50image→actions3 retries, 1 knob/try★ ≤ target8 · docsREADMEkernel_mismatch · INTEGRATION_PLAN± LATENCY_HANDBACKTOOL CATALOGUE (deterministic, hard-asserted)scan_checkpoint — inspect the saved model weights and their keysscan_reference_repo — read the reference code to find the model class and layer listwrite_spec_<hw> — declarative weight loader; checks: every weight loadswrite_pipeline_<hw> — the ~1,200-line forward; checks: finite outputs on dummy inputwrite_frontend_<hw> — buffer allocation + graph capture; checks: load_model worksregister_pipeline — registry + allow-list editsrun_correctness_gate — oracle + 20-shot run + alignment ≥ 0.99 vs referencemeasure_e2e_latency — wall-clock image→actions P50write_model_docs — integration README, plan, and noteswrite_latency_handoff — escalation doc to an engineerBLOCK: 3 consecutive correctness fails → INTEGRATION_BLOCKED.md, exit non-zeroSOFT-FAIL: 3 latency retries exhausted → LATENCY_HANDBACK.md, continue
The Model Integration Agent's internals. The orchestrator reads the 8-step workflow; each step is a sub-agent with deterministic tools and a hard gate. Gates 6 (alignment ≥ 0.99) and 7 (≤ target) are the load-bearing contracts.

How the design addresses the four difficulties:

  1. Silent failures → addressed by the gate combination in step 6: run the reference oracle, compare alignment, run 20 inferences with different noise seeds, check for NaN.
  2. Shape-dependent pieces → addressed by the scouted attention-spec template and the per-shape kernel-routing table built into step 3.
  3. GPU-graph capture → addressed by the pre-allocation contract in step 4 (every buffer the pipeline writes to is allocated in the frontend constructor, never inside the forward pass) and by the warmup-before-capture pattern enforced by step 5.
  4. Shape-and-hardware kernel selection → addressed by a per-(shape, hardware) routing table that the correctness and latency gates share. The agent picks a default kernel, runs the correctness gate, and if the alignment drop signals a kernel mismatch, switches to the next candidate and re-validates.

Three layers of guardrails, not one. The agent never trusts a single check.

  1. Tool-level assertions. Every tool returns structured data and asserts a specific invariant.
  2. Output-parser status. Every agent message ends with PASS or FAIL. The orchestrator refuses to advance on FAIL.
  3. Three-attempt retry budget. Each sub-agent has three attempts. After that, the orchestrator halts and writes a block report.

A worked example — adding SmolVLA, step by step

SmolVLA is a 32-layer VLA with a 12-layer vision tower, a 32-layer language model, and a 32-layer action expert with a 10-step denoising head. The target was our edge server. The agent's job was to produce a fully functional integration with both the no-NaN correctness gate and the end-to-end latency budget met.

The rough layout — step 1, scout. Before writing code, the scout sub-agent reads the reference model code and produces a layout that lists, per layer, which kernel to call and which buffer to write to. The layout also maps the differences between the desktop GPU and the server hardware:

AspectDesktop GPUServer / edge GPU
Attention backendFlash-Attention 2 (vendored)PyTorch SDPA with group-query-attention broadcasting
FP8 GEMMPer-tensor dynamicPer-tensor with precomputed descale, cuBLASLt-backed
Norm + activationBF16 split / FP8 quantize separatelyFused norm + FP8 quantize + GEMM
Rotation mathSingle fused split + rotate + cache-write kernelPlain split kernel + separate shape-aware rotation step
Buffer modelBackend-owned Q/K/VPipeline-owned Q/K/V + KV cache

The desktop and server pipelines cannot share code. They share a contract — pointer-only forward, fixed-shape buffers, contiguous inputs — but the kernel set, attention layout, and FP8 conventions all differ. The agent emits one file per hardware target, with no runtime forks. That dual-frontend rule exists exactly for this reason.

The plan and the implementation — steps 2–5. After the scout writes the integration plan, the agent emits a milestone plan with a hard gate at each of its eight checkpoints (M0 through M7), then works through steps 3–5. Those steps produce the bulk of the new code for an integration: a model-specific forward pass, a weight loader, and a runtime surface that allocates buffers, captures the GPU graph, and exposes the inference call. Three artifacts per integration, three substantial files. Two patterns recur in every new model:

  • Pre-allocated buffer pool. Every GPU buffer the pipeline writes to is allocated in the frontend constructor, never inside the forward pass — the hard contract for GPU-graph capture.
  • Forward as free functions. The pipeline class is a thin facade; the actual kernels are free functions that take raw pointers (vision_forward, text_decoder_forward, expert_forward).

Together, those patterns cover what integration actually is. The work is translating each block of the model — vision tower, language model, action expert with denoising head — into a sequence of kernel calls. This is where most of the per-model effort goes, and it is also the part the agent's templates handle best: the structural skeleton of a modern VLA is close enough across models that the templates generalize cleanly from one integration to the next.

The debug journey. The agent's correctness gate caught a sequence of issues, each pinned to a specific shape-and-kernel mismatch:

StageSymptomRoot causeFix
vision + connectoralignment 0.62 vs oracleWrong attention math (custom strided FMHA ≠ reference SDPA)Switch to torch SDPA + GQA broadcasting
text decoder KValignment 0.67 at layer 0 KRotation dim pairing wrong — fused kernel interleaves, reference does notUse a separate shape-aware rotation step
FP8 GEMMalignment 0.003 vs reference FP8 0.998Wrong FP8 kernel picked for this shape — CUTLASS path produces garbageSwitch to the cuBLASLt-backed FP8 path
FP8 descale alphafinal actions alignment 0.974–0.984Used act_scale + w_scale instead of productUse product α = a_s · w_s

Each row was a debugging cycle of 30 minutes to 2 hours. The goal is not to remove debugging entirely — it is to make sure each cycle is on a new failure, not a re-encounter of an old one. The orchestrator carries prior failures into its system prompt so the next integration skips the failure modes that have already been characterized.

Deliverables: latency and correctness check

Two things have to be true before the integration is shippable. Together they are the two deliverables of the agent.

Correctness check — no NaN, alignment ≥ 0.99 with the reference. The validation file the agent emits takes the new model and a reference PyTorch implementation, runs them on the same inputs, and compares the outputs. It then repeats that 20 times with different noise seeds to confirm the run is stable — and that no internal buffer ever produces a NaN, which is the canonical signal of a memory bug.

Two thresholds matter:

  • Internal alignment (text KV, intermediate v_t) ≥ 0.99 — proves the model structure is correct. Below 0.99 the integration is broken.
  • Final action alignment ≥ 0.75 in FP8 mode — proves the model produces usable actions. The number is lower than 0.99 because the 10-step denoising loop compounds per-step FP8 noise; the threshold reflects that.

For SmolVLA on the server: 20 of 20 inferences finite, text KV alignment 0.97–0.999 across all 32 layers, final actions alignment 0.78. If either threshold fails three times in a row, the agent writes a block report and stops — broken code is not committed.

End-to-end latency. Once correctness is green, the agent measures wall-clock image upload → set_prompt → graph replay → action download, as the median of 30 iterations after 10 warmups.

56.98 msP50 wall-clock (SmolVLA, server, incl. I/O)
47.54 msGraph replay only
0NaN across 20 inferences
0.998Alignment with reference (post-fix)
PathP50Note
SmolVLA on server (wall-clock)56.98 msIncl. image upload + sync; graph-replay only 47.54
SmolVLA on desktop Ada (FP8 + graph)~30 msDifferent hardware, already shipping
SmolVLA torch.compile (max-autotune)~62 msUser-stated baseline
SmolVLA TensorRT engine (BF16 batched)~42 msUser-stated baseline

The agent reports the best of five autotune trials — Thor tactic drift is non-deterministic (trial 0 → 49.10 ms, trial 4 → 47.46 ms). The ~9 ms gap between the graph-replay number and the wall-clock number is image upload (5 ms), noise upload (2 ms), and the final sync (1–2 ms).

Summary

The article has been about one shift: from treating model integration as a kernel-engineering problem to treating it as structured bookkeeping with a correctness + latency contract. The 1,200-line forward, the 200-line wrapper, the eight sub-agents and three guardrails are how that contract is enforced today. With SmolVLA, a single run produced a fully integrated, latency-budgeted pipeline in a fraction of the time the manual path used to take.

Next:

  • More models. Each new VLA trains the routing table and shortens the next integration.
  • More hardware. NPU and edge accelerators join the dual-frontend rule; per-hardware stays a hard rule.
  • Tighter loop. The Full SR Recovery Agent folds in, so correctness, latency and SR recovery live in one pass.

The point of the agent is not to remove engineers from the loop. It is to put them on the parts of the integration that the agent cannot honestly automate yet.