Model Deployment and Inference Optimization Questions

Situation: You need to deploy a neural network with tight latency and memory constraints on a target accelerator (GPU/TPU/NPU). Compiler toolchains - XLA, TVM, Glow - help by turning framework graphs into hardware-tailored code.

How they optimize (high-level comparison):

• Operator fusion: All three perform fusion to reduce kernel launches and memory traffic. XLA fuses at HLO-level for TensorFlow/TPU-first workloads; TVM has fine-grained graph- and tensor-level fusion with cost models; Glow does fusion oriented to CPU/accelerator backends with backend-specific node lowering.
• Layout transformation: XLA and Glow perform layout rewrites (e.g., NHWC↔NCHW) driven by backend conventions; TVM exposes explicit schedule primitives to change memory layout and tiling for vectorization and bandwidth utilization.
• Autotuning: TVM’s autotuner (AutoTVM/Ansor) searches schedules (tile sizes, unroll, parallelization) empirically on device; Glow relies more on static lowering plus backend-specific heuristics; XLA uses cost-model heuristics and some backend profiling (TPU benefits most).

Practical pipeline considerations:

• Profiling-first: run representative workloads to identify hotspots; use profilers (nvprof, perf, TPU profiler) before tuning.
• Incremental workflow: start with graph-level optimizations (pruning, op fusion), then layout transforms, then autotune kernels for hotspots.
• CI and reproducibility: store tuned schedules/artifacts (TVM’s log, XLA compile artifacts) in CI to avoid runtime autotune.
• Quantization & correctness: validate numeric equivalence and measure accuracy after quantization/fusion; include degradation budgets and fallback paths (reference CPU kernels).
• Deployment constraints: binary size, runtime JIT vs AOT (TVM supports both), target driver/runtime availability (Glow often used in embedded, XLA for TF/XLA runtime).
• DevOps: automate benchmarking, rollouts (canary), and add telemetry to detect regressions.

Example minimal flow:

1. Profile model in-framework to find heavy ops.
2. Export graph to toolchain (XLA HLO / TVM Relay / Glow IR).
3. Apply graph optimizations + layout rewrites.
4. Autotune kernels for target hardware (TVM’s Ansor or platform-specific tools).
5. AOT compile artifacts, run end-to-end validation and integrate into CI/CD.

Trade-offs:

• TVM: highest control and autotuning power, more engineering effort.
• XLA: best integrated with TF/TPU, good for server-side with less manual tuning.
• Glow: optimized for heterogeneous embedded contexts, simpler backend integration.

This approach balances performance, correctness, and maintainability for production ML deployments.

Approach summary:

• Build the custom CUDA kernel as an ONNX Runtime (ORT) custom op library (shared object / DLL) with a clear ABI and semantic versioning. Provide CPU fallback implementation in the same library or a separate one. Package per-platform artifacts (manylinux wheel for Python, Windows wheel, and a native deb/rpm/container image).
• At runtime, load and register the custom op library with ORT; detect GPU/driver capability and choose GPU kernel if available, otherwise use CPU fallback or alternative implementation.
• CI/CD builds cross-platform, runs unit+integration tests on CPU and GPU runners, signs and publishes artifacts with version metadata.

Concrete example (minimal ORT custom op skeleton + Python loader):

C++ custom op (simplified):

// custom_op.cc
#include "onnxruntime_cxx_api.h"

// MyKernel holds the actual (CUDA or CPU) implementation; it is a plain struct,
// not a subclass of any ORT-provided "OrtKernel" type.
struct MyKernel {
  MyKernel(const OrtApi& api, const OrtKernelInfo* info) { /* read attributes via api */ }
  void Compute(OrtKernelContext* context) { /* launch CUDA kernel or CPU fallback */ }
};

// Current ORT custom-op signature takes `const OrtApi&`, not the removed
// `Ort::CustomOpApi` wrapper (dropped from onnxruntime_cxx_api.h).
struct MyOp : Ort::CustomOpBase<MyOp, MyKernel> {
  const char* GetName() const noexcept { return "MyCustomOp"; }
  const char* GetExecutionProviderType() const noexcept { return "CUDAExecutionProvider"; }
  size_t GetInputTypeCount() const noexcept { return 1; }
  ONNXTensorElementDataType GetInputType(size_t) const noexcept { return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; }
  size_t GetOutputTypeCount() const noexcept { return 1; }
  ONNXTensorElementDataType GetOutputType(size_t) const noexcept { return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; }
  void* CreateKernel(const OrtApi& api, const OrtKernelInfo* info) const {
    return new MyKernel(api, info);
  }
};

extern "C" {
ORT_API_STATUS(ORT_API_CALL, RegisterCustomOps)(OrtSessionOptions* options, const OrtApiBase* api_base) {
  static Ort::CustomOpDomain domain("custom_domain");
  static MyOp op;
  domain.Add(&op);
  const OrtApi* ort_api = api_base->GetApi(ORT_API_VERSION);
  return ort_api->AddCustomOpDomain(options, domain);
}
}

Python loader and runtime selection:

import onnxruntime as ort
import torch
def make_session(onnx_path, lib_path):
    sess_opts = ort.SessionOptions()
    try:
        sess_opts.register_custom_ops_library(lib_path)
    except Exception as e:
        print("Custom op library load failed:", e)
    providers = []
    if ort.get_device() == 'GPU':
        providers = ['CUDAExecutionProvider','CPUExecutionProvider']
    else:
        providers = ['CPUExecutionProvider']
    return ort.InferenceSession(onnx_path, sess_options=sess_opts, providers=providers)

Packaging & versioning:

• Follow SemVer for the custom-op API, and pin against a specific ORT release: the custom-op C++ surface has changed across ORT versions (e.g., the old Ort::CustomOpApi wrapper type was removed from onnxruntime_cxx_api.h; CreateKernel now takes const OrtApi& directly, and a CreateKernelV2 status-returning variant also exists), so a kernel library built against one ORT ABI version is not guaranteed to load against another. Tag each release with git tag and include:
  • libcustomop.so / .pyd / .dll per platform
  • Python wheel manylinux2014_x86_64 / aarch64, Windows wheels; include package metadata exposing op version and required ORT version.
  • Docker images (gpu/cpu) with matching lib and driver compatibility matrix.

CI & tests:

• Matrix builds: ubuntu-latest (CPU), ubuntu+CUDA runner, windows-latest with CUDA, aarch64 where relevant.
• Steps: checkout, build native (CMake + nvcc), run unit tests (gtest) for kernel correctness, run integration tests: run ONNX model with known inputs and compare outputs to reference (PyTorch/NumPy) with tolerances.
• GPU tests validate device selection and numerical fidelity; include driver-check gating.
• Add fuzz/edge-case tests and performance benchmarks; run nightly to detect regressions.

Cross-platform build tips:

• Use CMake + CUDA toolkit, build static parts portable; use GitHub Actions with matrix; use manylinux builder (auditwheel) for Python wheels; use MSVC + vcpkg/CUDA on Windows.
• Containerize build environments to ensure reproducible builds; embed ABI info into artifact filenames.

Runtime versioning & compatibility:

• Export a function in the library reporting (op_version, build_cuda_version, ort_abi_hash).
• At load time, validate ORT version compatibility and CUDA runtime availability; if incompatible, fail fast and log guided error suggesting compatible builds.

Fallback strategies:

• Provide CPU implementation in the same custom-op library and register both kernels; at session creation set providers so ORT picks CUDA when available and CPU otherwise.
• If GPU exists but driver lacks feature, detect capability via cudaDriverGetVersion / cuInit and choose CPU path.
• Optionally implement a pure-ONNX equivalent graph as a fallback and ship it; guard with config flag to force CPU-only for debugging.

Observability & safety:

• Emit telemetry (build id, op version, selected kernel) at startup.
• Add health checks and a compatibility matrix; expose env vars to override kernel selection (e.g., CUSTOM_OP_PREFER_CPU=1).
• Use signed artifacts and CI gating (unit+integration + model-level E2E) before promotion.

Trade-offs:

• Single library with both implementations simplifies deployment but increases artifact size.
• Separate CPU/GPU artifacts reduce size but complicate packaging/selection.
• Prefer reproducible containerized builds and semantic versioning to minimize runtime surprises.

Approach:
Use a greedy, layer-wise search. For each layer i in order, attempt to lower its bitwidth from current (e.g., 8) to candidate smaller values (e.g., 4,2) in descending precision. Temporarily quantize that layer, evaluate on a validation metric; if metric degradation <= threshold, accept and lock the lower bitwidth and move to next layer. Otherwise revert and try next candidate. Return final per-layer bitwidths.

Interface and example implementation:

from typing import List, Callable, Sequence

def greedy_mixed_precision_search(
    model,                                # model object with apply_bitwidth(layer_idx, bits) and clone()/restore()
    val_fn: Callable[[object], float],    # validation function: returns metric (higher is better)
    layers: Sequence[int],                # list of layer indices [0..L-1]
    candidates: Sequence[int] = (8,4,2),  # bitwidth candidates, descending precision
    metric_drop_threshold: float = 0.01   # relative allowed drop (e.g., 0.01 => 1%)
) -> List[int]:
    """
    Returns list of selected bitwidths per layer.
    """
    # baseline metric with all layers at highest precision (candidates[0])
    for idx in layers:
        model.apply_bitwidth(idx, candidates[0])
    baseline = val_fn(model)

    selected = [candidates[0]] * len(layers)

    for i, layer in enumerate(layers):
        locked = selected[i]
        for bits in candidates[1:]:               # try lower precisions
            # apply candidate
            model.apply_bitwidth(layer, bits)
            metric = val_fn(model)
            # relative drop <= threshold?
            if (baseline - metric) / max(baseline, 1e-12) <= metric_drop_threshold:
                selected[i] = bits
                # update baseline to current metric to allow cumulative degradation tracking
                baseline = metric
                break
            else:
                # revert to previous locked precision
                model.apply_bitwidth(layer, locked)
        # proceed to next layer (locked choice)
    return selected

Key points / reasoning:

• Greedy reduces search from exponential to linear in L by locking layers after decision.
• Updating baseline to current metric makes threshold check cumulative; alternatively keep original baseline to bound total drop.

Complexity analysis:

• Let C = number of lower-bit candidates tried per layer (<= len(candidates)-1). Each attempt calls val_fn (cost V).
• Time: O(L * C * V). If V is expensive (full validation epoch), use a cheaper proxy (calibration dataset, subset).
• Space: O(1) extra aside from model; if model cloning used, space can be O(size(model)) temporarily.

Edge cases and enhancements:

• Non-monotonic metric changes: consider multiple passes (iterative greedy) to recover opportunities after earlier choices.
• Correlated layers: consider group-wise search or simulated annealing if greedy fails.
• If val_fn is noisy, average multiple runs or use statistical tests before locking.

I'll treat each technique with a concise description, how it reduces p99, and its trade-offs on cost, average latency, and worst-case latency.

Speculative decoding

• What: Run a small fast model (speculator) to predict several next tokens; verify in parallel with the large model that continues from the last verified token, or accept speculator outputs when validated.
• p99 impact: Reduces p99 substantially when speculator is usually correct; short-circuits long large-model runs.
• Cost: Increases total compute (two models sometimes run), but amortized if speculator avoids large-model steps; may raise cost modestly.
• Average latency: Often decreases (fast path); when speculator wrong, overhead increases.
• Worst-case latency: Slightly worse than baseline (extra validation) if speculator consistently fails.

Early-exit / prediction confidence

• What: At each step, use internal confidence scores or smaller early-exit heads to stop generation early or skip computation when confident.
• p99 impact: Lowers p99 for inputs the model is confident on; adaptive saving on long tails.
• Cost: Lowers cost overall if many tokens exit early; minimal extra model complexity.
• Average latency: Decreases.
• Worst-case latency: Unchanged or slightly better if confidence misestimates are rare; risk of quality loss if threshold badly tuned.

Hedged requests

• What: Send parallel requests to multiple replicas (same model or different sizes); use the first response.
• p99 impact: Directly reduces tail latency because you pick the fastest replica.
• Cost: Increases linearly with number of hedged replicas (higher compute & network).
• Average latency: Slightly increases average system load but client-perceived average latency falls.
• Worst-case latency: Greatly improved for tails but can still be high if all replicas slow or overloaded.

Prioritized scheduling

• What: Scheduler prioritizes short or latency-sensitive requests (e.g., small prompts, premium users) and preempts long jobs.
• p99 impact: Improves p99 for prioritized class; overall system tail may shift if starvation occurs.
• Cost: Neutral; better utilization.
• Average latency: Improves for prioritized traffic; may degrade for low-priority batch jobs.
• Worst-case latency: For high-priority traffic, worst-case drops; for deprioritized jobs, worst-case can increase.

Caching

• What: Cache token completions, logits, or final outputs for repeated prompts or prompt prefixes.
• p99 impact: Strongly reduces p99 for cache hits (near-zero latency).
• Cost: Lowers compute cost for hits but adds storage and cache management complexity.
• Average latency: Improves proportional to hit rate.
• Worst-case latency: Unchanged for cache misses; if cache hit rate unpredictable, tail remains.

Combined strategy

• Combine caching + spec decoding + prioritized scheduling for best practical p99 reduction: cache fast paths, use speculator to cover common prefixes, hedge only on truly latency-sensitive requests, and schedule prioritization to protect SLAs. Tune thresholds (confidence, hedge fanout, priority rules) and monitor quality/cost trade-offs.

Situation: We need a production inference architecture for a 100B-parameter LLM serving real‑time chat with ~50ms token latency. Design choices below prioritise lowest per‑token latency while keeping cost and GPU memory constraints realistic.

High-level approach

• Use mixed model-parallel strategy: intra-node tensor (operator) parallelism for GPU compute balance and inter-node pipeline or expert-stage split only if necessary. For a 100B model on modern GPUs (H100/A100 80GB), tensor parallelism (TP) + ZeRO-style optimizer state is sufficient for inference; pipeline parallelism adds scheduling overhead and extra micro-batching latency so avoid deep pipeline stages for 50ms targets.
• Place the full model weights sharded across GPUs with tensor parallel groups (Megatron-style) to keep each GPU compute-bound and minimise cross-node latency.

Model sharding

• Tensor parallelism (1D or 2.5D) for splitting weight matrices across GPUs; this preserves single-layer forward compute locality and allows fused kernels.
• Use model weight sharding (stateless) plus replication of small control state. Avoid multi-stage pipeline unless model > what fits in aggregate GPU memory with reasonable TP group sizes.

Inter‑GPU networking

• Prefer NVLink intra-node to minimize latency/frequency and use NCCL over NVLink for collectives. For multi-node setups use HDR InfiniBand with NCCL/RDMA enabled to get <5–10µs collective latency. Use 2-level NCCL: intra-node NVLink, inter-node NCCL over IB for reductions/broadcasts.

Activation checkpointing/offload

• For inference, activation checkpointing is less useful; instead:
  • Precompute and store static activations (if any) where possible.
  • Use activation recompute only for memory-limited edge cases - recompute increases latency, so avoid for 50ms target.
  • Use weight and KV cache offload: keep weights and hot KV cache on GPU; cold session KV can be paged to CPU RAM or NVMe via a fast async transfer path.

Memory and compute balancing

• Aim each GPU to hold its shard of weights + local attention compute + active KV cache. Use H100 80GB (or A100 80GB) nodes with 4–8 GPUs per node depending on chosen TP degree.
• Balance TP degree so per-GPU memory usage <= ~70% to leave headroom for KV caches and temporary buffers.
• Use quantized weights (4-bit or 8-bit) with high-quality quantization (GPTQ or AWQ) to reduce memory and increase throughput with TensorRT/fused kernels.

Batching strategies & latency control

• Use dynamic micro-batching (token-level batching): coalesce tokens arriving within a short window (e.g., 1–3ms) into a single forward. Maintain a strict max-wait scheduler: if coalescing would exceed 50ms SLA (service-level agreement), run immediately.
• Use request prioritization and per-session fairness. For chat, maintain session affinity so KV cache remains local to GPU; sticky routing reduces cache misses.
• Use beam size =1 for interactive latency unless necessary.

KV cache handling for chat sessions

• Keep per-session past key/value tensors on GPU as long as the session is active and within memory budget - this enables O(1) incremental decoding per token.
• For many concurrent sessions, hierarchical caching:
  • Hot set (recently active) resident on GPU RAM.
  • Warm set in host DRAM; asynchronously transfer to GPU when session becomes active.
  • Cold sessions serialized to NVMe, with prefetch on session wake.
• Use compressed storage (float16/quantized) for KV cache to reduce footprint. For long contexts, use attention window truncation + summarization to bound KV growth.

Software stack & kernels

• Use PyTorch + Triton/TensorRT fused attention and matmul kernels, NCCL for collectives. Use cuBLASLt for quantized GEMMs. Use a low-latency RPC (gRPC/Custom binary) and an inference scheduler (e.g., NVIDIA Triton or custom C++ server with fine-grained scheduling).

Hardware and cost trade-offs

• H100 (80GB) vs A100 (80GB): H100 gives higher FLOPS (lower latency) but is costlier. For strict 50ms target, H100 or A100 with TensorRT + quantization is recommended.
• More GPUs per replica (higher TP) reduces per-GPU memory but increases all-reduce collectives; choose TP degree that balances NVLink bandwidth and effective batch execution time.
• Single-node multi-GPU with NVLink is lowest-latency option; multi-node increases network dependency and jitter but scales capacity.
• Quantization and smaller precision reduce hardware needs and cost but require validation for accuracy/regression.

Trade-offs summary

• Minimize pipeline depth to reduce per-token scheduling latency; favour tensor parallelism to keep kernels fused and latency low.
• Keep KV cache on GPU for lowest latency; offload cold sessions to CPU/NVMe to scale concurrent users.
• Use aggressive quantization and optimized kernels to meet 50ms on fewer GPUs - this reduces cost but may slightly impact model quality.
• If SLA tightly enforces 50ms at high QPS (queries per second), invest in more powerful GPUs (H100) and NVLink-enabled multi-GPU nodes to avoid inter-node NCCL overhead.

Expected outcome

• With H100 80GB nodes, tensor-parallel shards across 4–8 GPUs, quantized weights, fused kernels, token-coalescing micro-batching, and GPU-resident KV caches for hot sessions, you can realistically achieve ~30–50ms per token for 100B LLM inference at modest concurrency.