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.
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:
- 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.
- 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.
- GPU-graph capture is unforgiving. No memory copies, no Python branches, no buffer allocation after the first call.
- 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.
How the design addresses the four difficulties:
- 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.
- Shape-dependent pieces → addressed by the scouted attention-spec template and the per-shape kernel-routing table built into step 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.
- 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.
- Tool-level assertions. Every tool returns structured data and asserts a specific invariant.
- Output-parser status. Every agent message ends with PASS or FAIL. The orchestrator refuses to advance on FAIL.
- 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:
| Aspect | Desktop GPU | Server / edge GPU |
|---|---|---|
| Attention backend | Flash-Attention 2 (vendored) | PyTorch SDPA with group-query-attention broadcasting |
| FP8 GEMM | Per-tensor dynamic | Per-tensor with precomputed descale, cuBLASLt-backed |
| Norm + activation | BF16 split / FP8 quantize separately | Fused norm + FP8 quantize + GEMM |
| Rotation math | Single fused split + rotate + cache-write kernel | Plain split kernel + separate shape-aware rotation step |
| Buffer model | Backend-owned Q/K/V | Pipeline-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:
| Stage | Symptom | Root cause | Fix |
|---|---|---|---|
| vision + connector | alignment 0.62 vs oracle | Wrong attention math (custom strided FMHA ≠ reference SDPA) | Switch to torch SDPA + GQA broadcasting |
| text decoder KV | alignment 0.67 at layer 0 K | Rotation dim pairing wrong — fused kernel interleaves, reference does not | Use a separate shape-aware rotation step |
| FP8 GEMM | alignment 0.003 vs reference FP8 0.998 | Wrong FP8 kernel picked for this shape — CUTLASS path produces garbage | Switch to the cuBLASLt-backed FP8 path |
| FP8 descale alpha | final actions alignment 0.974–0.984 | Used act_scale + w_scale instead of product | Use 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.
| Path | P50 | Note |
|---|---|---|
| SmolVLA on server (wall-clock) | 56.98 ms | Incl. image upload + sync; graph-replay only 47.54 |
| SmolVLA on desktop Ada (FP8 + graph) | ~30 ms | Different hardware, already shipping |
| SmolVLA torch.compile (max-autotune) | ~62 ms | User-stated baseline |
| SmolVLA TensorRT engine (BF16 batched) | ~42 ms | User-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.