Developer Tools Machine Learning 34 vues

cfregly/ai-performance-engineering: GPU Training & Inference Scaling Guide

B
Bright Coding
Auteur
cfregly/ai-performance-engineering: GPU Training & Inference Scaling Guide

Modern AI workloads routinely saturate GPUs without delivering proportional throughput. Engineers profile, tune, and still hit memory walls, communication bottlenecks, and kernel inefficiencies that erode cost-per-token economics. The gap between theoretical FLOPs and realized goodput—useful work completed per unit time—remains the central challenge in production AI systems.

cfregly/ai-performance-engineering addresses this directly. Maintained by Chris Fregly, a performance engineer with experience at Netflix, Databricks, and AWS↗ Bright Coding Blog, this repository provides code, labs, and resources tied to the O'Reilly book AI Systems Performance Engineering. It covers GPU optimization, distributed training, inference scaling, and full-stack tuning for modern AI workloads. With 1,662 stars, 233 forks, and active development through July 2026, it serves as a practical reference for engineers who need to move beyond benchmark chasing toward empirical, profile-first optimization.

[INTERNAL_LINK: PyTorch performance tuning]

What is cfregly/ai-performance-engineering?

cfregly/ai-performance-engineering is an open-source companion repository (Apache 2.0 licensed) to the O'Reilly book AI Systems Performance Engineering, published November 2025. The project sits at the intersection of systems engineering and machine learning infrastructure: it provides thousands of lines of PyTorch and CUDA C++ code examples targeting modern NVIDIA GPUs, alongside profiling data and case studies from production-scale workloads.

The maintainer, Chris Fregly, brings credibility through direct experience building AI/ML products at three major platform companies and authoring two prior O'Reilly books (Data Science on AWS and Generative AI on AWS). He also created the O'Reilly course "High-Performance AI in Production with NVIDIA GPUs" and co-produced the DeepLearning.ai course "Generative AI with Large-Language Models" with Andrew Ng.

The repository's relevance stems from timing: as models scale toward trillion parameters and clusters expand to multi-million GPU deployments, manual tuning becomes unsustainable. The repository documents compiler-driven acceleration (PyTorch compiler, OpenAI Triton), disaggregated inference architectures, and emerging AI-assisted optimization techniques—positioning it for engineers navigating this transition.

The repository is actively maintained (last commit July 6, 2026) and supported by a monthly meetup with 100,000+ members across 20+ cities, indicating sustained community engagement beyond typical book-accompanying projects.

Key Features

Profile-first methodology with production tooling. The repository emphasizes diagnostic workflows using Nsight Systems, Nsight Compute, and the PyTorch profiler—not surface-level metrics like GPU utilization, but stall-point analysis that reveals where cycles actually go. This aligns with the book's central tenet: optimize for goodput, not utilization.

Compiler and kernel-level optimization. Coverage spans the PyTorch compiler stack (torch.compile), OpenAI Triton for kernel generation without C++ boilerplate, and CUTLASS for peak-performance kernels. Advanced topics include inline PTX, SASS tuning, and cooperative thread block clusters—techniques typically scattered across NVIDIA documentation and conference talks.

Distributed training parallelism strategies. The repository documents data parallelism (DP), FSDP, tensor parallelism (TP), pipeline parallelism (PP), context parallelism (CP), and mixture-of-experts (MoE) configurations, with emphasis on overlapping computation and communication to minimize bubble overhead.

Inference serving at scale. Detailed coverage of vLLM, SGLang, TensorRT-LLM, and NVIDIA Dynamo includes disaggregated prefill/decode architectures, paged KV cache management, and dynamic routing strategies. This addresses the specific challenge of serving trillion-parameter models cost-effectively.

200+ item performance checklist. A field-tested appendix captures optimizations across the entire lifecycle: OS and driver tuning, NUMA awareness, GPU programming patterns, network optimization, power/thermal management, and reproducibility practices. This functions as a regression-prevention tool for teams.

AI-assisted optimization frontier. The repository and book explore emerging techniques including AlphaTensor-discovered algorithms, automated GPU kernel optimization, and reinforcement learning agents for runtime tuning—acknowledging that manual optimization will not scale to next-generation clusters.

Use Cases

Training throughput optimization for large language models. Engineering teams hitting memory bandwidth limits during multi-node training can apply the repository's NCCL tuning guidance, FSDP configuration patterns, and profiling workflows to identify and eliminate communication bubbles. The emphasis on overlapping computation with communication directly addresses the dominant cost driver in distributed training.

High-throughput inference serving. Platform teams deploying vLLM or TensorRT-LLM can implement disaggregated prefill/decode architectures documented in chapters 15-18, separating latency-sensitive decode phases from throughput-oriented prefill computation. The KV cache optimization strategies (including FlashMLA, ThunderMLA, and FlexDecoding) target the memory movement bottlenecks that dominate inference costs.

Kernel customization without C++ expertise. ML engineers needing custom operators can leverage the OpenAI Triton examples to generate GPU kernels through Python↗ Bright Coding Blog, avoiding the development velocity cost of CUDA C++ while achieving near-peak performance. The repository's Triton deep dive (Chapter 14) bridges from basic syntax to advanced implementations.

Production debugging and regression prevention. Teams can apply the 200+ item checklist during design reviews and incident post-mortems, using the structured categories (system architecture, GPU programming, inference serving, etc.) to ensure consistent coverage. The reproducibility and documentation practices address a common failure mode where optimizations work once but cannot be maintained.

Heterogeneous cluster optimization. Organizations with mixed GPU generations or accelerator types can apply the profiling methodologies to characterize actual performance per watt and throughput per dollar across configurations, rather than assuming nominal specifications translate linearly.

Installation & Setup

The repository follows standard GitHub conventions. Clone and explore the structure:

# Clone the repository
git clone https://github.com/cfregly/ai-performance-engineering.git
cd ai-performance-engineering

# Navigate to code examples
cd code/

The code/ directory contains organized examples corresponding to book chapters. Individual chapter directories contain Python and CUDA C++ sources with dependencies specified per-example. The repository does not use a monolithic package installation; instead, examples reference specific PyTorch, CUDA, and toolkit versions appropriate to each technique.

For the book's complete environment, consult the O'Reilly materials. The repository's CONTRIBUTING.md provides guidelines for code contributions, documentation improvements, and performance enhancements—indicating an expectation that users will engage actively rather than passively consume.

Key dependencies across examples include:

  • PyTorch (with torch.compile support)
  • NVIDIA CUDA Toolkit
  • Nsight Systems and Nsight Compute (for profiling chapters)
  • OpenAI Triton (for kernel customization chapters)
  • vLLM, SGLang, or TensorRT-LLM (for inference chapters)

The repository assumes Linux environments with NVIDIA GPUs; specific OS tuning guidance appears in Chapter 3 (OS, Docker↗ Bright Coding Blog, and Kubernetes Tuning).

Real Code Examples

The repository contains thousands of lines of PyTorch and CUDA C++ across its chapters. Below are representative patterns from the documented content.

Advertisement

PyTorch compilation for optimized execution:

import torch

# torch.compile integration for model optimization
# Applies graph-level optimizations, operator fusion, and backend selection
model = torch.compile(model, mode="max-autotune")

# The repository explores full compiler stack behavior including
# Triton backend code generation and XLA integration

This pattern from Chapter 13 demonstrates the entry point for PyTorch 2.x compiler optimization. The repository extends this with profiling integration, memory tuning, and distributed compilation strategies.

OpenAI Triton kernel definition:

import triton
import triton.language as tl

# Custom kernel without CUDA C++ boilerplate
# Triton abstracts block-level parallelism and memory coalescing
@triton.jit
def custom_kernel(input_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(axis=0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    x = tl.load(input_ptr + offsets, mask=mask)
    # Kernel computation here
    tl.store(output_ptr + offsets, result, mask=mask)

From Chapter 14, this illustrates Triton's Python-native approach to GPU kernel development. The repository progresses from this foundation to advanced implementations including fused attention variants and custom reduction operations.

Nsight Systems profiling annotation:

import torch.cuda.nvtx as nvtx

# Mark regions for Nsight Systems timeline analysis
# Critical for identifying stall points in distributed workloads
nvtx.range_push("forward_pass")
output = model(input)
nvtx.range_pop()

nvtx.range_push("all_reduce")
dist.all_reduce(tensor)
nvtx.range_pop()

This NVTX instrumentation pattern from Chapter 13 enables the profile-first methodology. The repository demonstrates correlating these markers with Nsight Systems timelines to distinguish computation from communication delays.

Note: The README references code organization but does not inline complete, runnable scripts in its rendered form. Engineers should explore the code/ directory directly for executable examples.

Advanced Usage & Best Practices

Adopt empirical measurement over assumption. The repository consistently emphasizes profiling before optimizing. Nsight Systems for timeline analysis and Nsight Compute for kernel-level inspection should precede any configuration change. This prevents the common anti-pattern of applying optimizations that improve microbenchmarks but hurt end-to-end throughput.

Structure optimizations around goodput, not utilization. A GPU at 100% utilization may still achieve minimal useful work if stalled on memory or synchronization. The repository's methodology targets reducing time-to-completion for actual model training steps or inference requests.

Leverage compiler stacks incrementally. Start with torch.compile in default mode, profile, then escalate to custom Triton kernels only for verified bottlenecks. The productivity cost of hand-written CUDA rarely justifies the performance gain unless standard tools prove insufficient.

Apply the checklist systematically. The 200+ item checklist in docs/appendix.md serves as both optimization guide and regression prevention tool. Consider integrating relevant sections into code review checklists for ML infrastructure changes.

Engage with the community ecosystem. The monthly meetups and YouTube channel provide ongoing content beyond the static repository, including conference recaps (NVIDIA GTC, NeurIPS) and vendor-specific deep dives that complement the book's foundational material.

Comparison with Alternatives

Aspect cfregly/ai-performance-engineering NVIDIA Deep Learning Performance Guide PyTorch Performance Tuning Guide
Scope Full-stack: hardware to inference GPU-centric, vendor-specific Framework-specific, narrower
Format Book + code repository + community Documentation + whitepapers Official docs + tutorials
Depth Production case studies with profiling data Reference specifications API-level guidance
Inference focus vLLM, SGLang, TensorRT-LLM, Dynamo TensorRT, Triton Inference Server Limited native coverage
Distributed training DP, FSDP, TP, PP, CP, MoE with overlap analysis NCCL optimization focus PyTorch Distributed APIs
Maintainership Individual expert + community NVIDIA engineering Meta/PyTorch team

cfregly/ai-performance-engineering distinguishes itself through integrated coverage spanning hardware architecture through serving infrastructure, with explicit cost-per-token economics. The NVIDIA guides offer deeper single-vendor specificity but less cross-stack integration. The PyTorch official documentation provides authoritative API reference but lacks the systems-level perspective and empirical case studies. Trade-off: this repository requires more synthesis effort; official guides offer more prescriptive, immediately applicable configuration snippets.

FAQ

What GPU hardware does this target? Modern NVIDIA GPUs including Grace CPU/Blackwell GPU configurations, with Tensor Core and Transformer Engine coverage.

Is this suitable for AMD or Intel accelerators? Primary focus is NVIDIA CUDA; some PyTorch patterns transfer but GPU-specific optimization targets NVIDIA architecture.

What PyTorch version is required? Examples use PyTorch 2.x with torch.compile; specific versions noted per-example in code/ directory.

Can I use this without purchasing the book? The repository contains substantial standalone code and resources; the book provides structured narrative and deeper context.

Is commercial use permitted under the license? Apache 2.0 allows commercial use with attribution requirements.

How current is the inference engine coverage? Includes vLLM, SGLang, TensorRT-LLM, and NVIDIA Dynamo as of 2025-2026; meetup content tracks emerging developments.

Does this cover quantization? Chapter 16 includes quantization approaches for real-time inference; NVFP4 and other low-precision numerics appear in meetup materials.

Conclusion

cfregly/ai-performance-engineering offers a rare combination: systems-level depth with practitioner-level accessibility, maintained by an author with direct experience at the platforms defining modern AI infrastructure. It is best suited for ML engineers, systems engineers, and platform teams who have moved past initial model deployment and now face the harder problem of sustainable, cost-efficient scale.

The repository does not promise easy wins. Its value lies in structured methodology—profile first, optimize empirically, validate end-to-end—and in comprehensive coverage of techniques from kernel-level CUDA through trillion-parameter inference architectures. For teams where GPU costs dominate operational expenditure, this represents a practical investment in engineering capability.

Explore the code, apply the checklist, and engage with the active community at the monthly meetups. Start at the source: https://github.com/cfregly/ai-performance-engineering.

Advertisement

Commentaires 0

Aucun commentaire pour l'instant. Soyez le premier à réagir !

Laisser un commentaire

Advertisement