Compiler

2. PyTorch → ONNX export

Trace the nn.Module with example inputs, lower aten ops to ONNX nodes, and pick an opset_version high enough for the ops you use. Resize cleaned up in opset 11; dynamic Slice in opset 10; opset 17 is a safe default for current PyTorch. Mark dynamic_axes for any batch or spatial dimension you want flexible. Python control flow and custom ops are the two recurring export hazards.

torch.onnx.export(
    model, x, "model.onnx",
    opset_version=17,
    dynamic_axes={"input": {0: "N"}, "output": {0: "N"}},
    do_constant_folding=True,
)

3. Graph optimization

Pre-build passes on the ONNX graph. Constant folding evaluates weight-only subgraphs at compile time. BatchNorm folds into the preceding Conv: the BN's affine collapses into Conv's weight and bias, so one fewer kernel ships. Identity and Dropout get removed. Layout transform rewrites the graph to NHWC when the NPU prefers channels-last. onnxsim or TensorRT's built-in optimizer does all of this.

python -m onnxsim model.onnx model.sim.onnx
# Conv(W, b) -> BN(gamma, beta, mean, var) -> ReLU
# becomes Conv(W', b') -> ReLU, where
#   s  = gamma / sqrt(var + eps)
#   W' = W * s
#   b' = (b - mean) * s + beta

4. INT8 post-training quantization (PTQ)

Feed a small calibration set (100–500 unlabeled images) through the FP32 graph, collect activation histograms per tensor, and pick a scale and zero-point that minimize KL divergence against the float distribution. Weights are quantized per-channel, one scale per output filter; activations per-tensor. INT8 holds up on conv stacks because per-channel weight scales and ReLU-bounded activations keep the dynamic range tight in each tensor. First conv and the final classifier are usually the sensitive layers; pin them to FP16 with a precision constraint.

# First pass: build a calibration cache from images in ./calib/
polygraphy convert model.sim.onnx --int8 \
    --data-loader-script calib_loader.py \
    --calibration-cache calib.cache \
    -o /dev/null

# Second pass: build the deployment engine, consuming the cache
trtexec --onnx=model.sim.onnx --int8 \
        --calib=calib.cache --saveEngine=model.int8.plan

5. Kernel fusion

The builder walks the graph and replaces matchable subgraphs with single fused kernels: Conv + Bias + ReLU, Conv + Add (residual), depthwise + pointwise (mobilenet block). Intermediates stay in registers / shared memory instead of round-tripping to LPDDR. The TensorRT verbose log shows [LayerFusion] lines listing which patterns matched.

6. TensorRT engine build

For each layer the builder times multiple tactics (different kernel implementations) on the actual hardware at the actual shape and precision, picks the winner, and serialises the choices into a .plan file. The plan is locked to one GPU / NPU SKU and one set of input shapes (unless you set a dynamic shape profile). This step is slow (minutes) but only happens once per deployment target.

7. Runtime on the edge SoC

Deserialize the plan, allocate I/O bindings in SoC SRAM where possible, set up a CUDA / vendor stream, and DMA the camera buffer straight into the input binding (zero-copy when the ISP shares the same physical memory). One enqueueV3 per frame; the NPU runs the fused kernel sequence; output comes back as detections / class scores.

// C++ pseudocode, edge inference loop
context->setTensorAddress("input",  d_input);
context->setTensorAddress("output", d_output);
while (capture(frame)) {
    dma(frame, d_input);
    context->enqueueV3(stream);
    cudaStreamSynchronize(stream);
    postprocess(d_output);
}

2. HF safetensors → GGUF or TensorRT-LLM

Two real paths. GGUF (llama.cpp): convert-hf-to-gguf.py model/ packs all weights + tokenizer + config into one file that the runtime mmaps; backend (CPU AVX / Metal / Vulkan / CUDA) is chosen at load time. TensorRT-LLM: a Python builder script defines the model in TRT-LLM's Python API, calls trtllm-build with model + max sequence / batch / KV-cache config, and emits an engine. ONNX is skipped: graph too dynamic, attention too custom.

# Path A: GGUF for CPU / NPU edge
python convert-hf-to-gguf.py meta-llama/Llama-3.2-1B --outfile model.gguf
llama-quantize model.gguf model.q4_k_m.gguf q4_k_m

# Path B: TensorRT-LLM for GPU edge
trtllm-build --checkpoint_dir ./hf_ckpt \
             --output_dir ./engine \
             --max_input_len 2048 --max_seq_len 4096 \
             --gemm_plugin float16 --use_paged_context_fmha enable

3. INT4 weight-only quantization

Decode is bandwidth-starved, so weights drop to INT4 or INT8 while activations stay FP16. GPTQ minimizes per-layer reconstruction error using calibration data and an approximate second-order Hessian (OBS-style). AWQ identifies the small fraction of weight channels that the activation outliers route through and scales them up before quantization, with an inverse scale on the activations so the math cancels out. Group size 64 or 128 sets the granularity (one scale per group). 4× fewer bytes per weight read from LPDDR. On Llama-class chat models at INT4 group 128, perplexity delta is typically under one point on common evals.

# AWQ via llm-awq, then export to TRT-LLM
python -m awq.entry --model_path meta-llama/Llama-3.2-1B \
                    --w_bit 4 --q_group_size 128 \
                    --run_awq --dump_awq awq.pt

4. KV cache layout

Each attention layer caches K and V tensors of shape [num_heads, head_dim] for every token generated so far. Total size is 2 × n_layers × n_heads × head_dim × seq_len × bytes_per_elem, often hundreds of MB at long context and the largest single LPDDR resident on the edge SoC. The compiler picks paged KV: allocate in fixed-size blocks, reference each layer's cache through a block table. The pattern originated in vLLM and is now the default in TRT-LLM. Sliding-window attention caps the cache at the last N tokens for models trained that way (Mistral-style).

5. Attention fusion

The attention block collapses into a small number of fused kernels. Q, K, and V projections fuse into one packed matmul against the concatenated weights. Fused multi-head attention (FlashAttention-style) computes softmax(QK^T)V in tiled passes without ever materializing the full S×S attention matrix. RMSNorm + matmul fuses so the normalization output stays in registers for the next projection. Rotary position applies inline to Q and K before the dot product. Each fusion replaces one LPDDR write-then-read pair with a register hop. On decode that translates directly into tokens/sec.

6. Engine / GGUF build

The builder produces the deployable artifact. With llama.cpp the GGUF file is the engine; the runtime ships the kernels and picks a backend at startup. With TRT-LLM, trtllm-build times tactics per layer at the chosen shapes, sets the KV cache config, and emits a serialized engine plus a runtime config JSON. Like CNN engines, a TRT-LLM plan is bound to one GPU SKU.

7. Prefill / decode loop on the edge SoC

Two regimes share one engine. Prefill runs one forward pass over the entire prompt, populating the KV cache. It is matmul-heavy and pegs the NPU. Decode generates one token at a time. Each step reads all the weights plus the existing KV cache and appends one row, which pegs LPDDR bandwidth and leaves the NPU mostly idle. Continuous batching is off at batch=1 edge, so the one knob that moves tokens/sec is memory bandwidth utilization. Stages 3 through 5 exist to lift it.

// llama.cpp pseudocode at the edge (current API)
llama_model  *m   = llama_model_load_from_file("model.q4_k_m.gguf", mparams);
llama_context*ctx = llama_init_from_model(m, cparams);

llama_batch batch = llama_batch_init(/*n_tokens*/ 512, /*embd*/ 0, /*n_seq_max*/ 1);

// prefill the prompt
for (int i = 0; i < n_prompt; ++i) {
    llama_batch_add(batch, prompt_tokens[i], /*pos*/ i,
                    /*seq_ids*/ {0}, /*logits*/ i == n_prompt - 1);
}
llama_decode(ctx, batch);

// decode one token at a time
while (n_generated < max_tokens) {
    int id = llama_sampler_sample(sampler, ctx, -1);
    if (llama_vocab_is_eog(llama_model_get_vocab(m), id)) break;
    llama_batch_clear(batch);
    llama_batch_add(batch, id, n_prompt + n_generated, {0}, true);
    llama_decode(ctx, batch);
    emit(id);
}