// 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.
| name | pytorch-kernel-optimization |
| description | 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. |
.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; choose modes:
"default" balanced, "reduce-overhead" for small batches/CUDA graphs, "max-autotune" for peak perf, "max-autotune-no-cudagraphs" when graphs undesirable.fullgraph=True for whole-graph capture; set dynamic=False when shapes are static.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).CUDAExtension (-O3, --use_fast_math, -arch=sm_80); enforce input checks in C++ bindings; expose kernels via PYBIND11_MODULE.torch.cuda.amp + GradScaler; mix dtypes per op if needed; leverage bfloat16 when supported.checkpoint, checkpoint_sequential); use memory-efficient attention via scaled_dot_product_attention; consider activation offloading (CPU swap) when memory-bound.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.fuse_conv_bn), apply quantization (quant.fuse_modules, prepare, convert), prune weights via torch.nn.utils.prune; ensure evaluation mode during quantization calibration.torch.cuda.CUDAGraph; warm up then capture forward/backward; reuse static input/output buffers; note torch.compile(mode=\"reduce-overhead\") can leverage graphs automatically.torch.profiler.profile with CPU/CUDA activities, schedules, and tensorboard_trace_handler; enable record_shapes, profile_memory, with_stack.prof.key_averages().table(sort_by=\"cuda_time_total\"); iterate on hotspots.torch.compile mode chosen for workload; warmup complete; performance measured post-compilation.This skill should be used when optimizing AMD GPU kernels on MI300 using the aiter project, including running op tests, benchmarking, iterating on kernel changes, and recording results in the kernel experiment database.
This skill should be used when reasoning about GPU architecture fundamentals to guide kernel optimization choices such as memory hierarchy usage, execution model mapping, block sizing, and latency-aware tuning across HIP, Triton, and PyTorch.
This skill should be used when writing or tuning HIP kernels on AMD/NVIDIA GPUs, covering memory coalescing, shared-memory tiling, bank conflict avoidance, warp primitives, occupancy, vectorization, async ops, loop unrolling, and profiling.
This skill should be used when optimizing kernels in this repo and needing to consult past optimization experiments, or when recording the current optimization iteration back into the kernel experiment database.
MI300/CDNA3 architecture guide for HIP/Triton optimization—MFMA variants, dual register files, data formats, sparsity, LDS/GWS, and best practices.
CDNA3/MI300 HIP programming insights—chiplet/cache model, Infinity Cache, memory coherency, matrix cores, sparsity, and best practices.