:::tip 🎮 Interactive Playground Visualize this concept: Try the vLLM Architecture demo on the EngineersOfAI Playground - no code required. :::

Model Compilation and Optimization

8ms to 2ms Without Changing the Model

The computer vision team has a problem. ResNet-50 inference runs at 8ms per image in PyTorch eager mode. The product SLA is 3ms. The model is correct - 76.1% top-1 accuracy on ImageNet, exactly what the business needs. The team does not want to change it. They need to make it faster without retraining.

The senior ML engineer opens NVIDIA's Nsight Systems profiler on the serving node. The trace reveals the problem immediately: ResNet-50's forward pass launches 247 separate CUDA kernels. Each kernel has overhead: argument validation, thread block scheduling, memory allocation checks. For small intermediate tensors, the kernel launch overhead is larger than the actual computation. The GPU spends more time launching work than doing work.

She runs TensorRT to compile the model. TensorRT analyzes the entire computation graph, fuses adjacent operations into single kernels (ReLU after convolution becomes one kernel, not two), selects optimal kernel implementations for the specific GPU architecture, and pre-plans memory layouts to minimize data movement. The compiled engine runs the same ResNet-50 in 1.9ms - a 4.2× speedup with no change to model weights or accuracy.

This is compiler-level optimization. The model's mathematical operations are identical. What changes is how those operations map to hardware: fewer kernel launches, better memory access patterns, optimized arithmetic through layer fusion. The difference between an unoptimized and compiled model is often the difference between missing and meeting your SLA.

---

Why This Exists - The Impedance Mismatch Between Frameworks and Hardware

PyTorch is designed for flexibility. Its eager execution model evaluates operations immediately and independently, making debugging intuitive and dynamic control flow possible. But this flexibility has a cost: PyTorch has no visibility into what operations will be called next, so it cannot optimize across operation boundaries.

Consider output = relu(batchnorm(conv(x, w))). In PyTorch eager mode:

1. conv(x, w) → launch Conv2d CUDA kernel, write 56×56×256 tensor to GPU memory
2. batchnorm(...) → read that tensor, launch BatchNorm kernel, write new tensor
3. relu(...) → read that tensor, launch ReLU kernel, write final tensor

Three separate memory round-trips through GPU DRAM. Three kernel launches. Three temporary tensors allocated and freed.

A compiler sees this as a dataflow graph and recognizes that the intermediate tensors never need to leave the GPU's on-chip cache. It fuses all three operations into a single kernel that reads x and w, computes conv → batchnorm → relu, and writes only the final output. One memory read/write. One kernel launch. The computation is identical; the hardware utilization is dramatically better.

---

Historical Context

TensorRT was launched by NVIDIA in 2016 as a deep learning inference optimizer and runtime for NVIDIA GPUs. Its kernel fusion, quantization, and architecture-specific kernel selection made it the production standard for NVIDIA GPU inference. TensorRT 8 (2021) added dynamic shapes and improved quantization workflows. TensorRT-LLM (2023) extended these optimizations to transformers at scale.

ONNX (Open Neural Network Exchange) was introduced by Microsoft and Facebook in 2017 as a common IR (intermediate representation) for ML models, enabling export from any framework and import into any runtime. ONNX became the universal "shipping format" for ML models - train in PyTorch, export to ONNX, optimize with TensorRT or ONNX Runtime.

XLA (Accelerated Linear Algebra) was developed inside Google for TensorFlow in 2017 and extended to JAX. XLA performs whole-program compilation: it takes a computation graph and produces optimized machine code that runs efficiently on GPUs and TPUs, with automatic operator fusion, memory layout optimization, and SIMD vectorization.

torch.compile (PyTorch 2.0, 2023) brought compiler-level optimizations to the PyTorch ecosystem without requiring ONNX export. It uses torch.fx for graph capture, Triton for GPU kernel generation, and TorchInductor as the backend to generate optimized CUDA code. For many models, torch.compile achieves 1.5-3× speedup with a single line of code.

---

The Computation Graph and Why It Matters

Every neural network is a computation graph: nodes are operations (convolution, matrix multiply, activation), edges are tensors. Compilation works on this graph.

The performance model for fused operations: the bottleneck shifts from memory bandwidth to arithmetic throughput. Conv+BN+ReLU fused reads the input once, computes all three operations on-chip, writes the output once. Unfused, it reads and writes three times the data.

$\text{Latency}_{\text{unfused}} = \sum_i \left(\frac{D_i}{\text{BW}} + \frac{F_i}{\text{TFLOPS}}\right)$

$\text{Latency}_{\text{fused}} = \frac{D_{\text{in}} + D_{\text{out}}}{\text{BW}} + \frac{\sum_i F_i}{\text{TFLOPS}}$

Where $D_i$ is the data volume for operation $i$ and $F_i$ is its FLOPs. For operations where $D_i \gg F_i$ (memory-bound), fusion gives the most benefit.

---

torch.compile: One Line of Code, Real Speedup

torch.compile is the easiest entry point to compiled inference for PyTorch models:

# torch_compile_demo.py - applying compilation to a production model
import torch
import torchvision.models as models
import time
import numpy as np

device = torch.device("cuda")

# Load model
model = models.resnet50(weights='IMAGENET1K_V2').to(device)
model.eval()

# Compile - three modes with different tradeoffs
# 'default': good balance of compile time vs speedup
# 'reduce-overhead': more aggressive, best for repeated small batches
# 'max-autotune': exhaustive kernel search, slow compile, best runtime
compiled_model = torch.compile(model, mode='reduce-overhead')

def benchmark(fn, input_tensor, warmup=50, iters=200):
    """Accurate GPU benchmark with CUDA event timing."""
    # Warmup: let JIT compilation and GPU caches settle
    for _ in range(warmup):
        with torch.no_grad():
            fn(input_tensor)
    torch.cuda.synchronize()

    # Measure
    times = []
    start_event = torch.cuda.Event(enable_timing=True)
    end_event = torch.cuda.Event(enable_timing=True)

    for _ in range(iters):
        start_event.record()
        with torch.no_grad():
            fn(input_tensor)
        end_event.record()
        torch.cuda.synchronize()
        times.append(start_event.elapsed_time(end_event))

    return {
        'mean': np.mean(times),
        'p50': np.percentile(times, 50),
        'p95': np.percentile(times, 95),
        'p99': np.percentile(times, 99),
    }

# Benchmark batch of 32
x = torch.randn(32, 3, 224, 224, device=device)

eager_stats = benchmark(model, x)
compiled_stats = benchmark(compiled_model, x)

print(f"Eager:    mean={eager_stats['mean']:.2f}ms, p99={eager_stats['p99']:.2f}ms")
print(f"Compiled: mean={compiled_stats['mean']:.2f}ms, p99={compiled_stats['p99']:.2f}ms")
print(f"Speedup:  {eager_stats['mean'] / compiled_stats['mean']:.2f}×")
# Typical output on A100:
# Eager:    mean=12.4ms, p99=14.1ms
# Compiled: mean=5.8ms, p99=6.3ms
# Speedup:  2.14×

Handling Dynamic Shapes

By default, torch.compile recompiles when input shapes change. For serving where request sizes vary, use dynamic=True:

# For variable-length NLP inputs
compiled_model = torch.compile(model, dynamic=True)
# Or mark specific dimensions as dynamic:
from torch._dynamo import mark_dynamic

def predict_variable_length(model, input_ids):
    # Mark the sequence length dimension as dynamic
    mark_dynamic(input_ids, 1)  # dim 1 is sequence length
    with torch.no_grad():
        return model(input_ids)

---

TensorRT: Maximum Performance on NVIDIA GPUs

TensorRT is the production choice when you need every last millisecond. It performs more aggressive optimization than torch.compile: architecture-specific kernel selection, INT8/FP8 quantization built-in, and explicit memory planning.

Export and Optimize Flow

# tensorrt_export.py - convert PyTorch model to TensorRT engine
import torch
import torch.onnx
import tensorrt as trt
import numpy as np
from pathlib import Path

def export_to_onnx(
    model: torch.nn.Module,
    input_shape: tuple,
    onnx_path: str,
    dynamic_axes: dict = None
):
    """Export PyTorch model to ONNX format."""
    model.eval()
    dummy_input = torch.randn(*input_shape, device='cuda')

    torch.onnx.export(
        model,
        dummy_input,
        onnx_path,
        export_params=True,
        opset_version=17,           # use latest supported opset
        do_constant_folding=True,   # fold constant operations
        input_names=['input'],
        output_names=['output'],
        dynamic_axes=dynamic_axes or {
            'input': {0: 'batch_size'},
            'output': {0: 'batch_size'}
        }
    )
    print(f"ONNX model saved to {onnx_path}")


def build_tensorrt_engine(
    onnx_path: str,
    engine_path: str,
    fp16: bool = True,
    int8: bool = False,
    calibrator = None,
    max_batch_size: int = 64,
    workspace_gb: float = 4.0
) -> bytes:
    """
    Build TensorRT engine from ONNX model.
    This takes 2-10 minutes but only runs once per deployment.
    """
    logger = trt.Logger(trt.Logger.WARNING)
    builder = trt.Builder(logger)
    network = builder.create_network(
        1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
    )
    parser = trt.OnnxParser(network, logger)
    config = builder.create_builder_config()

    # Set workspace memory limit
    config.set_memory_pool_limit(
        trt.MemoryPoolType.WORKSPACE,
        int(workspace_gb * (1 << 30))
    )

    # Enable FP16 if GPU supports it
    if fp16 and builder.platform_has_fast_fp16:
        config.set_flag(trt.BuilderFlag.FP16)
        print("FP16 enabled")

    # Enable INT8 with calibrator
    if int8 and calibrator:
        config.set_flag(trt.BuilderFlag.INT8)
        config.int8_calibrator = calibrator
        print("INT8 enabled with calibration")

    # Parse ONNX
    with open(onnx_path, 'rb') as f:
        if not parser.parse(f.read()):
            for error in range(parser.num_errors):
                print(parser.get_error(error))
            raise RuntimeError("Failed to parse ONNX model")

    # Configure dynamic shapes (for variable batch size)
    profile = builder.create_optimization_profile()
    profile.set_shape(
        'input',
        min=(1, 3, 224, 224),        # minimum batch size
        opt=(32, 3, 224, 224),       # typical/optimal batch size
        max=(max_batch_size, 3, 224, 224)  # maximum batch size
    )
    config.add_optimization_profile(profile)

    print("Building TensorRT engine (may take several minutes)...")
    serialized_engine = builder.build_serialized_network(network, config)

    if serialized_engine is None:
        raise RuntimeError("Failed to build TensorRT engine")

    # Save engine to disk
    with open(engine_path, 'wb') as f:
        f.write(serialized_engine)
    print(f"TensorRT engine saved to {engine_path}")

    return serialized_engine


class TRTInference:
    """Load and run TensorRT engine for inference."""

    def __init__(self, engine_path: str):
        logger = trt.Logger(trt.Logger.WARNING)
        runtime = trt.Runtime(logger)

        with open(engine_path, 'rb') as f:
            self.engine = runtime.deserialize_cuda_engine(f.read())

        self.context = self.engine.create_execution_context()
        import pycuda.driver as cuda
        import pycuda.autoinit
        self.cuda = cuda

        # Allocate buffers once
        self.buffers = self._allocate_buffers()

    def _allocate_buffers(self):
        inputs, outputs, bindings = [], [], []
        stream = self.cuda.Stream()

        for i in range(self.engine.num_io_tensors):
            name = self.engine.get_tensor_name(i)
            dtype = trt.nptype(self.engine.get_tensor_dtype(name))
            shape = self.engine.get_tensor_shape(name)
            size = abs(trt.volume(shape))  # account for dynamic dims
            host_mem = self.cuda.pagelocked_empty(size, dtype)
            device_mem = self.cuda.mem_alloc(host_mem.nbytes)
            bindings.append(int(device_mem))

            if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
                inputs.append({'name': name, 'host': host_mem, 'device': device_mem})
            else:
                outputs.append({'name': name, 'host': host_mem, 'device': device_mem})

        return {'inputs': inputs, 'outputs': outputs,
                'bindings': bindings, 'stream': stream}

    def infer(self, input_array: np.ndarray) -> np.ndarray:
        """Run inference. input_array shape: (batch, 3, 224, 224)."""
        buf = self.buffers
        batch_size = input_array.shape[0]

        # Set dynamic shape
        self.context.set_input_shape('input', input_array.shape)

        # Copy input to device
        np.copyto(buf['inputs'][0]['host'], input_array.ravel())
        self.cuda.memcpy_htod_async(
            buf['inputs'][0]['device'],
            buf['inputs'][0]['host'],
            buf['stream']
        )

        # Execute
        self.context.execute_async_v3(stream_handle=buf['stream'].handle)

        # Copy output to host
        self.cuda.memcpy_dtoh_async(
            buf['outputs'][0]['host'],
            buf['outputs'][0]['device'],
            buf['stream']
        )
        buf['stream'].synchronize()

        return buf['outputs'][0]['host'].reshape(batch_size, -1)

---

ONNX Runtime: Portable Optimization

ONNX Runtime is the production inference runtime from Microsoft. It supports GPU (via CUDA/TensorRT execution providers) and CPU (via OpenMP, MLAS optimizations) with the same model file.

# onnxruntime_serving.py
import onnxruntime as ort
import numpy as np

def create_ort_session(
    onnx_path: str,
    use_gpu: bool = True,
    use_tensorrt: bool = False
) -> ort.InferenceSession:
    """Create an optimized ONNX Runtime session."""
    sess_options = ort.SessionOptions()
    sess_options.graph_optimization_level = (
        ort.GraphOptimizationLevel.ORT_ENABLE_ALL
    )
    sess_options.execution_mode = ort.ExecutionMode.ORT_PARALLEL
    sess_options.inter_op_num_threads = 4
    sess_options.intra_op_num_threads = 4

    if use_tensorrt and use_gpu:
        # Use TensorRT for GPU execution (maximum performance)
        providers = [
            ('TensorrtExecutionProvider', {
                'device_id': 0,
                'trt_max_workspace_size': 4 * (1 << 30),  # 4GB
                'trt_fp16_enable': True,
                'trt_engine_cache_enable': True,
                'trt_engine_cache_path': '/tmp/trt_cache',
            }),
            ('CUDAExecutionProvider', {'device_id': 0}),
            'CPUExecutionProvider',
        ]
    elif use_gpu:
        providers = [
            ('CUDAExecutionProvider', {
                'device_id': 0,
                'arena_extend_strategy': 'kNextPowerOfTwo',
                'gpu_mem_limit': 8 * (1 << 30),  # 8GB
                'cudnn_conv_algo_search': 'EXHAUSTIVE',
            }),
            'CPUExecutionProvider',
        ]
    else:
        providers = ['CPUExecutionProvider']

    session = ort.InferenceSession(
        onnx_path,
        sess_options=sess_options,
        providers=providers
    )
    print(f"Execution providers: {session.get_providers()}")
    return session


def batch_predict(
    session: ort.InferenceSession,
    images: np.ndarray  # shape: (N, 3, 224, 224)
) -> np.ndarray:
    input_name = session.get_inputs()[0].name
    output_name = session.get_outputs()[0].name

    results = session.run(
        [output_name],
        {input_name: images.astype(np.float32)}
    )
    return results[0]

---

Profiling: Finding the Real Bottlenecks

Before optimizing, profile. The Nsight Compute profiler and PyTorch's built-in profiler reveal where time actually goes.

# profiling.py - finding inference bottlenecks
import torch
from torch.profiler import profile, record_function, ProfilerActivity

def profile_model_inference(model, input_tensor, num_steps=50):
    """
    Profile model inference to identify bottlenecks.
    Captures CPU + CUDA activity with kernel details.
    """
    model.eval()

    with profile(
        activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
        record_shapes=True,
        profile_memory=True,
        with_stack=True
    ) as prof:
        for step in range(num_steps):
            with record_function(f"inference_step_{step}"):
                with torch.no_grad():
                    output = model(input_tensor)
            if step < 5:
                prof.step()

    # Print top operations by CUDA time
    print("=== Top 20 ops by CUDA time ===")
    print(prof.key_averages().table(
        sort_by="cuda_time_total",
        row_limit=20
    ))

    # Export for TensorBoard visualization
    prof.export_chrome_trace("trace.json")
    # Open in Chrome: chrome://tracing → Load → trace.json

    # Find the most expensive kernel
    key_avgs = prof.key_averages()
    top_op = max(key_avgs, key=lambda x: x.cuda_time_total)
    print(f"\nSlowest operation: {top_op.key}")
    print(f"  CUDA time: {top_op.cuda_time_total / 1e3:.2f}ms total "
          f"({top_op.cuda_time_total / top_op.count / 1e3:.3f}ms avg)")
    print(f"  Called {top_op.count} times")
    print(f"  Memory: {top_op.cpu_memory_usage / 1e6:.1f}MB CPU, "
          f"{top_op.cuda_memory_usage / 1e6:.1f}MB CUDA")


# Common findings from profiling:
# 1. Many small kernels → kernel fusion opportunity (torch.compile helps)
# 2. Large memory allocation/free → pre-allocate fixed buffers
# 3. Host-device copy overhead → pin memory, async transfers
# 4. Synchronization points → async execution with CUDA streams
# 5. Contiguous memory calls → ensure input tensors are contiguous

Memory Layout Optimization

Tensor memory layout significantly affects performance. The channels_first vs channels_last format can make a 10-30% difference for convolutions:

# memory_layout.py - channels_last for better convolution performance
import torch
import torchvision.models as models

model = models.resnet50().cuda()
model.eval()

# Convert model to use channels_last memory format
# (matches GPU's preferred memory layout for convolutions)
model = model.to(memory_format=torch.channels_last)

x = torch.randn(32, 3, 224, 224, device='cuda')
# Convert input to channels_last as well
x = x.to(memory_format=torch.channels_last)

# Benchmark - channels_last typically 10-20% faster for CNNs
with torch.no_grad():
    output = model(x)
    print(f"Output shape: {output.shape}")
    print(f"Output is contiguous: {output.is_contiguous()}")
    print(f"Output memory format: {output.is_contiguous(memory_format=torch.channels_last)}")

---

XLA for TPU and Multi-Device

XLA compiles computation graphs for Google TPUs and also supports GPUs via JAX. For models that run on Google Cloud TPUs, XLA compilation is mandatory for performance:

# jax_xla_example.py - JAX with XLA compilation
import jax
import jax.numpy as jnp
from jax import jit, vmap
import flax.linen as nn

class ResNetBlock(nn.Module):
    features: int

    @nn.compact
    def __call__(self, x):
        residual = x
        x = nn.Conv(self.features, (3, 3))(x)
        x = nn.BatchNorm(use_running_average=True)(x)
        x = nn.relu(x)
        x = nn.Conv(self.features, (3, 3))(x)
        x = nn.BatchNorm(use_running_average=True)(x)
        return nn.relu(x + residual)


# XLA compilation happens via @jit decorator
@jit
def predict_batch(params, images):
    """XLA-compiled batch prediction. First call compiles, subsequent calls are fast."""
    return model.apply(params, images, train=False)


# Vectorize over batch dimension automatically
batched_predict = vmap(predict_batch, in_axes=(None, 0))

# After first call (compilation), subsequent calls run optimized XLA program
x = jnp.ones((32, 224, 224, 3))
output = batched_predict(params, x)  # First call: compile + run (~5s)
output = batched_predict(params, x)  # Second call: just run (~1ms)

---

Compilation Comparison

Tool	Best For	Compile Time	Runtime Speedup	Dynamic Shapes
torch.compile	PyTorch models, easy adoption	30s-2min	1.5-3×	Yes (dynamic=True)
TensorRT	NVIDIA GPU, max performance	2-10min	2-8×	Yes (profiles)
ONNX Runtime + TRT	Portable, production	1-5min	2-6×	Limited
XLA/JAX	TPU, research, Google infra	5-30min	2-10× (TPU)	Limited
TorchScript	Export, mobile	10s	1.2-1.5×	No

---

Production Engineering Notes

When to Compile vs When Not To

Compile when:

• Model runs the same computation repeatedly (production serving)
• Batch size and input shapes are fixed or bounded
• You have time to run the compilation step in your CI/CD pipeline
• Model runs on GPU (compilation benefits are largest there)

Do not compile when:

• You are debugging model behavior (compiled code is harder to inspect)
• Dynamic control flow is central to the model (compilation handles it poorly)
• Input shapes change dramatically per request (recompilation costs)
• Prototype/research phase where model changes frequently

Engine Caching in CI/CD

TensorRT engines are device-specific. An engine built for A100 will not run on T4. Build and cache engines as part of your deployment pipeline:

# engine_manager.py - build and cache TRT engines
import hashlib
import os
from pathlib import Path

def get_or_build_engine(
    onnx_path: str,
    cache_dir: str,
    gpu_arch: str = None,
    **build_kwargs
) -> str:
    """
    Check cache for pre-built engine; build if missing.
    Cache key includes model hash + GPU architecture + build config.
    """
    if gpu_arch is None:
        import subprocess
        result = subprocess.run(
            ['nvidia-smi', '--query-gpu=name', '--format=csv,noheader'],
            capture_output=True, text=True
        )
        gpu_arch = result.stdout.strip().replace(' ', '_')

    # Compute cache key
    with open(onnx_path, 'rb') as f:
        model_hash = hashlib.sha256(f.read()).hexdigest()[:16]

    config_str = f"{gpu_arch}_{build_kwargs.get('fp16', False)}_{build_kwargs.get('int8', False)}"
    cache_key = f"{model_hash}_{config_str}"
    engine_path = Path(cache_dir) / f"engine_{cache_key}.trt"

    if engine_path.exists():
        print(f"Loading cached engine: {engine_path}")
        return str(engine_path)

    print(f"Building engine (no cache found for {cache_key})...")
    build_tensorrt_engine(onnx_path, str(engine_path), **build_kwargs)
    return str(engine_path)

---

Common Mistakes

:::danger Running TensorRT Benchmark Without Warmup TensorRT engines have a JIT compilation step on the first inference call and a GPU warm-up period. Benchmarking without 50-100 warmup iterations dramatically underestimates performance. The first call may be 5-10× slower than steady-state. Always warm up, always use CUDA event timing (not time.time()), and always synchronize before measuring. :::

:::danger Building TRT Engine on a Different GPU Than You Deploy To TensorRT engines are device-specific. An engine built for A100 Ampere may not load on T4 Turing, and vice versa. Build engines in your CI/CD pipeline on hardware that matches your production environment. Use the trt_engine_cache_enable option in ONNX Runtime to cache per-device. :::

:::warning ONNX Export Fails for Dynamic Control Flow torch.onnx.export traces the model with a specific input. Any Python-level control flow (if statements, loops) that depends on tensor values at runtime will be baked into the ONNX graph using the tracing input's values, silently producing an incorrect export. Use torch.jit.script for models with dynamic control flow before ONNX export, or check the ONNX model output carefully against the PyTorch model. :::

:::warning Compilation Overhead in Cold Starts torch.compile compiles on first call, causing a 30-second to 2-minute delay. In serverless or auto-scaling deployments, newly started instances pay this compilation cost on the first real request. Use compilation caching (TORCHINDUCTOR_CACHE_DIR) to persist compiled kernels across restarts, and warm up instances with synthetic requests before routing live traffic. :::

---

Interview Q&A

Q1: What is kernel fusion and why does it improve inference performance?

A: Kernel fusion combines multiple sequential GPU operations into a single kernel. The performance benefit comes from two sources. First, memory bandwidth: when operations are separate, intermediate tensors must be written to GPU DRAM and read back by the next operation. For a Conv+BatchNorm+ReLU sequence, this means 3 DRAM reads and 3 writes for the intermediate tensors. Fused, the intermediate results stay in the GPU's L2 cache or registers - 1 DRAM read and 1 write. Since transformer inference is largely memory-bandwidth-bound (90% of time is loading weights, not computing with them), reducing memory traffic directly reduces latency. Second, kernel launch overhead: each CUDA kernel has ~5-10 microseconds of scheduling overhead. A fused kernel replaces three launches with one. At 247 kernel launches for ResNet-50, that is roughly 1.5ms of pure overhead on a fast GPU. TensorRT and torch.compile both perform fusion automatically on the computation graph.

Q2: What is the difference between torch.compile and TensorRT, and when would you choose each?

A: torch.compile is the simpler option - one line of code, stays in the PyTorch ecosystem, handles dynamic shapes with dynamic=True, and produces 1.5-3× speedup for most models. It uses TorchInductor to generate Triton kernels specifically tuned for your operations. TensorRT requires an ONNX export step, GPU-specific engine compilation (2-10 minutes), and device-specific engines. But it performs more aggressive optimizations: architecture-specific kernel selection (choosing the best cublas/cudnn algorithm for your exact batch size and tensor shape), built-in INT8/FP8 calibration, and deeper cross-layer fusion. TensorRT typically achieves 2-8× speedup, higher than torch.compile. Choose torch.compile for simplicity, dynamic-shape workloads, and research workflows. Choose TensorRT when you need maximum throughput on NVIDIA hardware, are deploying to a fixed-shape production endpoint, and have the CI/CD infrastructure to manage engine builds.

Q3: How do you profile an inference bottleneck in PyTorch?

A: Start with PyTorch's built-in profiler (torch.profiler.profile with ProfilerActivity.CUDA). Export the trace to Chrome format and view it in chrome://tracing - this shows the timeline of every CUDA kernel, including duration and gaps. The profiler's key_averages().table(sort_by="cuda_time_total") output immediately shows the top time-consuming operations. Common findings: (a) many tiny kernels with high launch overhead → kernel fusion with torch.compile; (b) large host-to-device copies → use pinned memory (pin_memory=True in DataLoader) and async transfers; (c) synchronization points (torch.cuda.synchronize() called explicitly or by .item()) → remove or push to end; (d) convolution choosing a slow algorithm → set torch.backends.cudnn.benchmark = True to auto-tune. For deeper kernel-level analysis, NVIDIA Nsight Compute gives per-warp occupancy, memory bandwidth, and arithmetic throughput, identifying whether a kernel is memory-bound or compute-bound.

Q4: What is the role of ONNX in the ML deployment stack?

A: ONNX (Open Neural Network Exchange) is an intermediate representation format for ML models. Its role is decoupling: you can train in any framework (PyTorch, TensorFlow, JAX) and deploy via any runtime (ONNX Runtime, TensorRT, TVM, CoreML). The workflow is: train in PyTorch → export to ONNX (torch.onnx.export) → optimize and run via ONNX Runtime or TensorRT ONNX parser. ONNX Runtime with CUDA and TensorRT execution providers achieves near-TensorRT performance while supporting models from any source framework. The limitation is that ONNX covers a specific set of operators (opsets), and custom PyTorch operations may not export cleanly - you need to register custom ops or restructure the model. ONNX is also not ideal for models with significant Python-level dynamic control flow, as the export traces a specific execution path.

Q5: How does memory layout (channels_first vs channels_last) affect convolution performance?

A: GPUs process data in chunks that map to hardware memory access patterns. CUDA's cuDNN convolution kernels have different optimal memory layouts depending on the GPU architecture and operation type. Channels-last (NHWC: batch, height, width, channels) often outperforms channels-first (NCHW) for standard convolutions on Ampere and later GPUs, because it aligns with how the tensor cores access data for matrix multiplications. The difference is typically 10-30% for ResNet-class architectures. PyTorch's to(memory_format=torch.channels_last) converts the model's parameter layout to channels_last, and you must also convert inputs. torch.compile and TensorRT automatically select the optimal memory layout as part of compilation - one more reason to use compilation for production. Note: not all operations support channels_last; some will silently fall back to NCHW with an implicit copy, negating the benefit. Check for contiguous operations after conversion.