skill_example/references/a100-optimization-guide.md

A100 GPU Optimization Guide for CUDA Kernels

Overview

The NVIDIA A100 GPU is based on the Ampere architecture (sm_80) and remains one of the most widely deployed GPUs for AI workloads. This guide covers optimization strategies specific to A100 hardware when developing CUDA kernels for the HuggingFace Kernels ecosystem.

A100 Hardware Specifications

Specification	Value
Architecture	Ampere (sm_80)
Streaming Multiprocessors (SMs)	108
HBM2e Bandwidth	2.0 TB/s
Shared Memory per SM	164 KB (configurable)
L2 Cache	40 MB
FP32 CUDA Cores	6912
Tensor Cores (3rd gen)	432
Memory	40 GB or 80 GB HBM2e
TDP	400W (SXM) / 300W (PCIe)
Max Threads per SM	2048
Max Threads per Block	1024
Warp Size	32
Max Warps per SM	64

Shared Memory Configuration

The A100 supports flexible shared memory and L1 cache partitioning. The total 192 KB per SM can be split as:

Shared Memory	L1 Cache
0 KB	192 KB
28 KB	164 KB
100 KB	92 KB
132 KB	60 KB
164 KB	28 KB

For most kernel workloads in diffusers and transformers, 164 KB shared memory is optimal:

// Request maximum shared memory for a kernel
cudaFuncSetAttribute(
    my_kernel,
    cudaFuncAttributeMaxDynamicSharedMemorySize,
    164 * 1024  // 164 KB
);

Vectorized Memory Access

Vectorized loads and stores are critical for saturating the A100's 2.0 TB/s memory bandwidth. Always prefer wider data types for memory operations.

BF16 Vectorized Access

#include <cuda_bf16.h>

// Load 4 BF16 values at once (8 bytes) using float2
__device__ __forceinline__ void vectorized_load_bf16(
    const __nv_bfloat16* input,
    int idx,
    __nv_bfloat16& v0,
    __nv_bfloat16& v1,
    __nv_bfloat16& v2,
    __nv_bfloat16& v3
) {
    float2 packed = reinterpret_cast<const float2*>(input)[idx / 4];
    __nv_bfloat162 low = reinterpret_cast<__nv_bfloat162*>(&packed)[0];
    __nv_bfloat162 high = reinterpret_cast<__nv_bfloat162*>(&packed)[1];
    v0 = __low2bfloat16(low);
    v1 = __high2bfloat16(low);
    v2 = __low2bfloat16(high);
    v3 = __high2bfloat16(high);
}

FP16 Vectorized Access

#include <cuda_fp16.h>

// Load 8 FP16 values at once (16 bytes) using float4
__device__ __forceinline__ void vectorized_load_fp16_x8(
    const half* input,
    int idx,
    half vals[8]
) {
    float4 packed = reinterpret_cast<const float4*>(input)[idx / 8];
    half2* pairs = reinterpret_cast<half2*>(&packed);
    #pragma unroll
    for (int i = 0; i < 4; i++) {
        vals[2 * i] = __low2half(pairs[i]);
        vals[2 * i + 1] = __high2half(pairs[i]);
    }
}

FP32 Vectorized Access

// Load 4 FP32 values at once (16 bytes)
__device__ __forceinline__ float4 vectorized_load_fp32(
    const float* input,
    int idx
) {
    return reinterpret_cast<const float4*>(input)[idx / 4];
}

Alignment requirements: Vectorized loads require the base pointer to be aligned to the access width. For float4, this means 16-byte alignment. PyTorch tensors are typically 256-byte aligned, so this is usually satisfied.

Occupancy Tuning for 108 SMs

The A100 has 108 SMs, which means your grid dimensions must be tuned to this number for optimal utilization.

Grid Sizing Rules

// Rule 1: Grid should be a multiple of 108 for full utilization
// Rule 2: At minimum, launch 108 blocks (1 per SM)
// Rule 3: For best latency hiding, launch 2-4 blocks per SM

int num_sms = 108;
int blocks_per_sm = 2;  // Good default for compute-bound kernels
int grid_size = num_sms * blocks_per_sm;  // 216

// For memory-bound kernels, more blocks help hide latency
int blocks_per_sm_membound = 4;
int grid_size_membound = num_sms * blocks_per_sm_membound;  // 432

Occupancy Calculation

#include <cuda_runtime.h>

void print_occupancy(void* kernel, int block_size, int shared_mem) {
    int max_active_blocks;
    cudaOccupancyMaxActiveBlocksPerMultiprocessor(
        &max_active_blocks,
        kernel,
        block_size,
        shared_mem
    );

    int max_threads_per_sm = 2048;
    float occupancy = (float)(max_active_blocks * block_size) / max_threads_per_sm;
    printf("Block size: %d, Active blocks/SM: %d, Occupancy: %.1f%%\n",
           block_size, max_active_blocks, occupancy * 100);
}

Recommended Block Sizes

Kernel Type	Block Size	Blocks/SM	Occupancy
Element-wise (RoPE, activations)	256	8	100%
Row-wise reduction (LayerNorm, RMSNorm)	256	4-8	50-100%
Tiled matmul (attention)	128	4-8	25-50%
Memory-bound simple	512	4	100%

TF32 Mode

The A100 introduced TF32 (TensorFloat-32), which provides FP32-like range with reduced precision. TF32 uses 19 bits (1 sign + 8 exponent + 10 mantissa) versus FP32's 32 bits.

Enabling TF32

// TF32 is used automatically by Tensor Cores for FP32 inputs
// To explicitly control it:

// Enable TF32 for matmul (default on since CUDA 11.0 on A100)
at::globalContext().setAllowTF32CuBLAS(true);

// Enable TF32 for cuDNN convolutions
at::globalContext().setAllowTF32CuDNN(true);

In PyTorch

import torch

# Enable TF32 (default on A100)
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

# Disable for full FP32 precision when needed
torch.backends.cuda.matmul.allow_tf32 = False

When to Use TF32

• Use TF32: Training, inference where slight precision loss is acceptable, attention score computation
• Avoid TF32: Loss computation, gradient accumulation, numerical validation, reduction operations where precision matters

build.toml Configuration

Configure your kernel package for A100 targets:

[build]
cuda-version = "12.4"
cuda-capabilities = ["8.0"]

[kernel.rmsnorm]
src = ["rmsnorm.cu"]

[kernel.rope]
src = ["rope.cu"]

[kernel.gelu]
src = ["gelu.cu"]

Multi-GPU Target Configuration

To support both A100 and other GPUs:

[build]
cuda-version = "12.4"
# A100 (sm_80), A10G (sm_86), H100 (sm_90)
cuda-capabilities = ["8.0", "8.6", "9.0"]

Important: Each additional capability increases build time and binary size. Only include targets you actually need to support.

Compiler Flags for A100

[build]
cuda-version = "12.4"
cuda-capabilities = ["8.0"]
# A100-specific optimizations
extra-cuda-flags = [
    "--use_fast_math",
    "--maxrregcount=128",
    "-lineinfo"
]

Migration Tips from H100

If you have kernels optimized for H100, here are the key adjustments for A100:

Architecture Differences

Feature	H100 (sm_90)	A100 (sm_80)
SMs	132	108
HBM Bandwidth	3.35 TB/s	2.0 TB/s
Shared Memory/SM	228 KB (max)	164 KB (max)
L2 Cache	50 MB	40 MB
Tensor Core Gen	4th	3rd
TMA (Tensor Memory Accelerator)	Yes	No
Thread Block Clusters	Yes	No
FP8 Support	Yes	No

Code Changes Required

1. Remove TMA Usage

// H100 code using TMA -- NOT available on A100
// Replace with standard global memory loads

// Before (H100):
// cute::copy(tma_load, gmem_tensor, smem_tensor);

// After (A100):
// Use standard vectorized loads
float4 data = reinterpret_cast<const float4*>(global_ptr)[idx];

2. Adjust Grid Dimensions

// H100: 132 SMs
// A100: 108 SMs
// Always query at runtime:
int num_sms;
cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, 0);

3. Reduce Shared Memory Usage

// H100 allows up to 228 KB shared memory per block
// A100 allows up to 164 KB shared memory per block
// Reduce tile sizes or use multi-stage pipelining with smaller buffers

// H100 version: large tiles
// constexpr int TILE_SIZE = 256;  // May require > 164 KB

// A100 version: smaller tiles
constexpr int TILE_SIZE = 128;  // Fits within 164 KB

4. Remove Thread Block Clusters

// H100 supports cooperative launch with clusters
// A100 does not -- use standard cooperative groups

// Before (H100):
// dim3 cluster_size(2, 1, 1);
// cudaLaunchKernelEx(&config, kernel, args...);

// After (A100):
kernel<<<grid, block, shared_mem>>>(args...);

5. No FP8 -- Use BF16 or FP16

// H100 supports FP8 (E4M3 and E5M2)
// A100 does not -- use BF16 as the closest alternative

// Before (H100):
// __nv_fp8_e4m3 val = ...;

// After (A100):
__nv_bfloat16 val = __float2bfloat16(float_val);

A100-Specific Optimization Patterns

Async Memory Copy (cp.async)

The A100 supports asynchronous memory copies from global to shared memory, which is valuable for pipelining:

// Use cp.async to overlap memory transfers with computation
__device__ void async_load_to_shared(
    void* smem_ptr,
    const void* gmem_ptr,
    int bytes
) {
    asm volatile(
        "cp.async.ca.shared.global [%0], [%1], %2;\n"
        :: "r"(static_cast<unsigned>(__cvta_generic_to_shared(smem_ptr))),
           "l"(gmem_ptr),
           "n"(bytes)
    );
}

// Commit and wait for async copies
__device__ void async_commit_and_wait() {
    asm volatile("cp.async.commit_group;\n");
    asm volatile("cp.async.wait_group 0;\n");
    __syncthreads();
}

L2 Cache Residency Control

The A100's 40 MB L2 cache supports persistence hints:

// Set L2 cache persistence for frequently accessed data
cudaStreamAttrValue stream_attr;
stream_attr.accessPolicyWindow.base_ptr = (void*)data_ptr;
stream_attr.accessPolicyWindow.num_bytes = num_bytes;
stream_attr.accessPolicyWindow.hitRatio = 1.0f;
stream_attr.accessPolicyWindow.hitProp = cudaAccessPropertyPersisting;
stream_attr.accessPolicyWindow.missProp = cudaAccessPropertyStreaming;

cudaStreamSetAttribute(stream, cudaStreamAttributeAccessPolicyWindow, &stream_attr);

Profiling on A100

Using nsys

# Profile a Python script using your kernel
nsys profile --stats=true \
    --trace=cuda,cudnn,cublas,nvtx \
    -o a100_profile \
    python benchmark.py

Using ncu (Nsight Compute)

# Detailed kernel analysis
ncu --set full \
    --target-processes all \
    --launch-skip 5 --launch-count 10 \
    python benchmark.py

# Check memory throughput specifically
ncu --metrics \
    dram__bytes_read.sum,dram__bytes_write.sum,\
    l2__read_throughput.avg.pct_of_peak_sustained,\
    sm__throughput.avg.pct_of_peak_sustained \
    python benchmark.py

Key Metrics to Monitor

Metric	Good Value	Action if Low
DRAM Throughput	> 70% of 2.0 TB/s	Improve coalescing, use vectorized loads
SM Occupancy	> 50%	Reduce register usage, adjust block size
L2 Hit Rate	> 60%	Improve data locality, use cache hints
Shared Memory Efficiency	> 80%	Avoid bank conflicts
Warp Execution Efficiency	> 90%	Reduce branch divergence

Example: RMSNorm Kernel Optimized for A100

#include <cuda_bf16.h>
#include <cuda_runtime.h>

constexpr int WARP_SIZE = 32;
constexpr int A100_MAX_BLOCK = 1024;

template<int BLOCK_SIZE>
__global__ void rmsnorm_a100(
    const __nv_bfloat16* __restrict__ input,
    const __nv_bfloat16* __restrict__ weight,
    __nv_bfloat16* __restrict__ output,
    const int hidden_size,
    const float epsilon
) {
    const int row = blockIdx.x;
    const int tid = threadIdx.x;

    const __nv_bfloat16* row_input = input + row * hidden_size;
    __nv_bfloat16* row_output = output + row * hidden_size;

    // Phase 1: Compute sum of squares with vectorized loads
    float sum_sq = 0.0f;

    // Use float4 loads (8 BF16 values per load)
    const int vec_size = 8;
    const int num_vecs = hidden_size / vec_size;

    for (int i = tid; i < num_vecs; i += BLOCK_SIZE) {
        float4 packed = reinterpret_cast<const float4*>(row_input)[i];
        __nv_bfloat162* pairs = reinterpret_cast<__nv_bfloat162*>(&packed);
        #pragma unroll
        for (int j = 0; j < 4; j++) {
            float v0 = __bfloat162float(__low2bfloat16(pairs[j]));
            float v1 = __bfloat162float(__high2bfloat16(pairs[j]));
            sum_sq += v0 * v0 + v1 * v1;
        }
    }

    // Phase 2: Warp-level reduction
    #pragma unroll
    for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) {
        sum_sq += __shfl_xor_sync(0xffffffff, sum_sq, offset);
    }

    // Phase 3: Block-level reduction via shared memory
    __shared__ float warp_sums[BLOCK_SIZE / WARP_SIZE];
    int warp_id = tid / WARP_SIZE;
    int lane_id = tid % WARP_SIZE;

    if (lane_id == 0) {
        warp_sums[warp_id] = sum_sq;
    }
    __syncthreads();

    if (warp_id == 0) {
        sum_sq = (lane_id < BLOCK_SIZE / WARP_SIZE) ? warp_sums[lane_id] : 0.0f;
        #pragma unroll
        for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) {
            sum_sq += __shfl_xor_sync(0xffffffff, sum_sq, offset);
        }
    }

    __shared__ float rms_scale;
    if (tid == 0) {
        rms_scale = rsqrtf(sum_sq / hidden_size + epsilon);
    }
    __syncthreads();

    // Phase 4: Apply normalization with vectorized stores
    for (int i = tid; i < num_vecs; i += BLOCK_SIZE) {
        float4 in_packed = reinterpret_cast<const float4*>(row_input)[i];
        float4 w_packed = reinterpret_cast<const float4*>(weight)[i];

        __nv_bfloat162* in_pairs = reinterpret_cast<__nv_bfloat162*>(&in_packed);
        __nv_bfloat162* w_pairs = reinterpret_cast<__nv_bfloat162*>(&w_packed);

        float4 out_packed;
        __nv_bfloat162* out_pairs = reinterpret_cast<__nv_bfloat162*>(&out_packed);

        #pragma unroll
        for (int j = 0; j < 4; j++) {
            float v0 = __bfloat162float(__low2bfloat16(in_pairs[j]));
            float v1 = __bfloat162float(__high2bfloat16(in_pairs[j]));
            float w0 = __bfloat162float(__low2bfloat16(w_pairs[j]));
            float w1 = __bfloat162float(__high2bfloat16(w_pairs[j]));

            out_pairs[j] = __halves2bfloat162(
                __float2bfloat16(v0 * rms_scale * w0),
                __float2bfloat16(v1 * rms_scale * w1)
            );
        }

        reinterpret_cast<float4*>(row_output)[i] = out_packed;
    }
}

Summary

When targeting the A100:

1. Use sm_80 in your build configuration
2. Tune grids for 108 SMs -- aim for multiples of 108 blocks
3. Leverage 164 KB shared memory with appropriate L1/shared partitioning
4. Vectorize all memory access to approach the 2.0 TB/s bandwidth ceiling
5. Use TF32 for matrix operations where full FP32 precision is not required
6. Profile with ncu to verify you are hitting bandwidth and compute targets
7. When migrating from H100, remove TMA, clusters, FP8, and reduce shared memory tile sizes