# Pytorch Kernel Optimization

> This skill should be used when optimizing PyTorch models and kernels, including efficient tensor operations, torch.compile, custom autograd/CUDA/Triton extensions, mixed precision, memory and data pipeline tuning, model optimization techniques, CUDA graphs, and profiling.

- Skill: `wenyi-li/pytorch-kernel-optimization` (Agent Skill)
- Install (CLI): `npx skillmds@latest add wenyi-li/pytorch-kernel-optimization`
- Raw SKILL.md: https://api.skillmd.com/api/skills/wenyi-li/pytorch-kernel-optimization/raw
- Safety review: pending
- Works with: Claude Code, Claude.ai, OpenAI Codex
- Category: DevOps & Infra
- Author: wenyi-li (https://skillmd.com/u/wenyi-li)
- Updated: 2026-09-22
- Page: https://skillmd.com/skills/wenyi-li/pytorch-kernel-optimization

---


# PyTorch Kernel Optimization

## Purpose
- Equip PyTorch workflows with concrete optimization patterns from high-level APIs to custom kernels.
- Provide practical snippets for compilation, extensions, mixed precision, memory efficiency, and profiling.

## When to Use
- Tuning PyTorch models for throughput/latency on GPU.
- Deciding between compiler-level optimizations and custom kernels (C++/CUDA/Triton).
- Profiling and addressing bottlenecks in compute or input pipelines.

## How to Use
- **Efficient tensor ops**: favor contiguous layouts (`.contiguous()` when needed); use `channels_last` for convs; replace Python loops with vectorized ops; prefer in-place ops (`add_`, `mul_`, `out=`) when autograd-safe.
- **torch.compile**: wrap functions or models with `@torch.compile`; choose modes:
  - `"default"` balanced, `"reduce-overhead"` for small batches/CUDA graphs, `"max-autotune"` for peak perf, `"max-autotune-no-cudagraphs"` when graphs undesirable.
  - Use `fullgraph=True` for whole-graph capture; set `dynamic=False` when shapes are static.
- **Custom autograd**: implement `torch.autograd.Function` saving minimal tensors; recompute in backward when memory-bound (e.g., checkpointed attention); use custom backward formulas for fused ops (e.g., SiLU).
- **CUDA extensions**: build with `CUDAExtension` (`-O3`, `--use_fast_math`, `-arch=sm_80`); enforce input checks in C++ bindings; expose kernels via `PYBIND11_MODULE`.
- **Mixed precision**: train with `torch.cuda.amp` + `GradScaler`; mix dtypes per op if needed; leverage `bfloat16` when supported.
- **Memory optimization**: apply gradient checkpointing (`checkpoint`, `checkpoint_sequential`); use memory-efficient attention via `scaled_dot_product_attention`; consider activation offloading (CPU swap) when memory-bound.
- **Data loading**: configure `DataLoader` with `num_workers`, `pin_memory`, `prefetch_factor`, `persistent_workers`, `drop_last`; implement fast collate; prefetch to GPU with custom loader using streams and non-blocking copies.
- **Model optimization**: fuse Conv+BN (`fuse_conv_bn`), apply quantization (`quant.fuse_modules`, `prepare`, `convert`), prune weights via `torch.nn.utils.prune`; ensure evaluation mode during quantization calibration.
- **CUDA graphs**: capture steady workloads via `torch.cuda.CUDAGraph`; warm up then capture forward/backward; reuse static input/output buffers; note `torch.compile(mode=\"reduce-overhead\")` can leverage graphs automatically.
- **Profiling**:
  - Use `torch.profiler.profile` with CPU/CUDA activities, schedules, and `tensorboard_trace_handler`; enable `record_shapes`, `profile_memory`, `with_stack`.
  - Review `prof.key_averages().table(sort_by=\"cuda_time_total\")`; iterate on hotspots.

## Validation Checklist
- Tensor layouts contiguous/channels_last as appropriate; Python loops eliminated; in-place ops safe for autograd.
- `torch.compile` mode chosen for workload; warmup complete; performance measured post-compilation.
- Custom ops (autograd or CUDA) validate device/contiguity; register usage and block sizes tuned for kernels.
- AMP scaling stable (no inf/nan); dtype choices align with numerical sensitivity.
- Data loader keeps GPU fed (no data starvation); streams overlap transfers where applicable.
- Profiling reviewed after each major change; bottlenecks addressed or noted.