菜单
Optimize dense matrix multiplication (GEMM) kernels in Triton for NVIDIA and AMD GPUs. Covers tiled blocking, L2 cache grouping, tensor core utilization, and dual-platform autotune. Use when writing or optimizing matmul, linear layers, or any GEMM-based kernel.
用 Codex 或 Claude 帮你安装 复制这段 Prompt，粘贴到 Codex、Claude 或其他助手里，让它检查 Skill 页面并帮你完成安装。
基于 SOC 职业分类
| name | gemm-kernel-optimization |
| description | Optimize dense matrix multiplication (GEMM) kernels in Triton for NVIDIA and AMD GPUs. Covers tiled blocking, L2 cache grouping, tensor core utilization, and dual-platform autotune. Use when writing or optimizing matmul, linear layers, or any GEMM-based kernel. |
General Matrix Multiply (C = A @ B) is the most compute-intensive operation in deep learning. GEMM is compute-bound: the key metric is TFLOPS (target: >60% of hardware peak).
| Bound | Metric | Target (H100) | Target (4090) | Target (MI300X) |
|---|---|---|---|---|
| Compute | TFLOPS | >500 FP16 | >150 FP16 | >600 FP16 |
[BLOCK_M, BLOCK_N] tiles, one per program instanceBLOCK_K chunks[BLOCK_M, BLOCK_K], B tile [BLOCK_K, BLOCK_N], call tl.dotNaive row-major block ordering causes excessive HBM reads. Group GROUP_SIZE_M rows of blocks
before moving to the next column to maximize L2 reuse of B tiles.
Empirically: GROUP_SIZE_M=8 gives 10-15% improvement on A100/H100.
tl.dot(a, b, acc) maps to WMMA (Ampere) or WGMMA (Hopper) or MFMA (CDNA)NVIDIA_CONFIGS = [
triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=3, num_warps=8),
triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=5, num_warps=2),
]
AMD_CONFIGS = [
triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_warps=4, num_stages=2),
triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_warps=8, num_stages=2),
triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_warps=8, num_stages=2),
triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8, 'waves_per_eu': 8}, num_warps=4, num_stages=2),
]
Key AMD differences:
num_warps=4 means 256 threadsnum_stages must be >= 1 (0 crashes on HIP)waves_per_eu hint for occupancy tuning (2-4 compute-bound, 4-8 memory-bound)python skills/kernels/gemm/test_gemm.py
torch.allclose(triton_result, torch.matmul(A, B), atol=1e-2, rtol=1e-2) for FP162*M*N*K / latency_s / 1e12.to(output_dtype) before store: silent precision issuesWhen optimizing a GEMM kernel:
triton_template.py as baselinetest_gemm.py after every change to verify correctness + measure TFLOPSOptimize fused cross-entropy loss kernels in Triton for NVIDIA and AMD GPUs. Covers fused log-softmax + NLL, online log-sum-exp, and large vocabulary handling. Use when writing or optimizing cross-entropy, label-smoothed CE, or similar loss kernels.
Optimize FlashAttention-style fused attention kernels in Triton for NVIDIA and AMD GPUs. Covers online softmax, tiled QK/AV GEMM, causal masking, and memory-efficient attention. Use when writing or optimizing self-attention, cross-attention, or any QKV attention kernel.
Optimize RMS Normalization kernels in Triton for NVIDIA and AMD GPUs. Covers row-parallel reduction, vectorized loads, warp shuffle patterns, and rsqrt usage. Use when writing or optimizing RMSNorm, LayerNorm, or similar normalization kernels.
Optimize Rotary Position Embedding (RoPE) kernels in Triton for NVIDIA and AMD GPUs. Covers interleaved/sequential layouts, sincos precomputation, and multi-row processing. Use when writing or optimizing RoPE, position embeddings, or rotary transformations.
Optimize fused softmax kernels in Triton for NVIDIA and AMD GPUs. Covers online max/sum reduction, multi-row processing, numerical stability, and warp-level patterns. Use when writing or optimizing softmax, log-softmax, or similar row-wise normalization kernels.
Chooses block/tile dimensions, warp counts, pipeline stages, and L2-aware block grouping for GPU kernels (Triton/CUDA/HIP). Use when autotuning matmul-like kernels, tuning occupancy, or aligning tile geometry with M/N/K and NVIDIA vs AMD thread models.