# 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:

```cpp
// 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

```cpp
#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

```cpp
#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

```cpp
// 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

```cpp
// 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

```cpp
#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

```cpp
// 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

```python
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:

```toml
[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:

```toml
[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

```toml
[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**

```cpp
// 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**

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

**3. Reduce Shared Memory Usage**

```cpp
// 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**

```cpp
// 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**

```cpp
// 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:

```cpp
// 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:

```cpp
// 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

```bash
# 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)

```bash
# 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

```cpp
#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