Split View: CUDA GPU 프로그래밍 심화: Warp 최적화, Tensor Core, Triton 커널 작성까지

// Tiled Matrix Multiplication with Shared Memory
#define TILE_SIZE 16

__global__ void tiledMatMul(
    float* A, float* B, float* C,
    int M, int N, int K
) {
    __shared__ float tileA[TILE_SIZE][TILE_SIZE];
    __shared__ float tileB[TILE_SIZE][TILE_SIZE];

    int row = blockIdx.y * TILE_SIZE + threadIdx.y;
    int col = blockIdx.x * TILE_SIZE + threadIdx.x;
    float sum = 0.0f;

    int numTiles = (K + TILE_SIZE - 1) / TILE_SIZE;

    for (int t = 0; t < numTiles; t++) {
        // 타일을 shared memory에 로드
        int aCol = t * TILE_SIZE + threadIdx.x;
        int bRow = t * TILE_SIZE + threadIdx.y;

        tileA[threadIdx.y][threadIdx.x] =
            (row < M && aCol < K) ? A[row * K + aCol] : 0.0f;
        tileB[threadIdx.y][threadIdx.x] =
            (bRow < K && col < N) ? B[bRow * N + col] : 0.0f;

        __syncthreads();

        // 타일 내에서 행렬 곱 수행
        for (int k = 0; k < TILE_SIZE; k++) {
            sum += tileA[threadIdx.y][k] * tileB[k][threadIdx.x];
        }

        __syncthreads();
    }

    if (row < M && col < N) {
        C[row * N + col] = sum;
    }
}

__constant__ float convWeights[256];

__global__ void convKernel(float* input, float* output, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        float result = 0.0f;
        for (int i = 0; i < 256; i++) {
            result += input[idx + i] * convWeights[i]; // 브로드캐스트 접근
        }
        output[idx] = result;
    }
}

// 나쁜 예: Warp divergence 발생
__global__ void badKernel(float* data, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx % 2 == 0) {          // 워프 내 분기 발생
        data[idx] = data[idx] * 2.0f;
    } else {
        data[idx] = data[idx] + 1.0f;
    }
}

// 좋은 예: 워프 단위로 분기 정렬
__global__ void goodKernel(float* data, int n) {
    int warpId = (blockIdx.x * blockDim.x + threadIdx.x) / 32;
    int idx    = blockIdx.x * blockDim.x + threadIdx.x;
    if (warpId % 2 == 0) {       // 워프 단위 분기 — divergence 없음
        data[idx] = data[idx] * 2.0f;
    } else {
        data[idx] = data[idx] + 1.0f;
    }
}

// 나쁜 예: Column-major 접근 (stride = N)
__global__ void columnAccess(float* matrix, int M, int N) {
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    float sum = 0.0f;
    for (int row = 0; row < M; row++) {
        sum += matrix[row * N + col];   // 스트라이드 N 접근
    }
}

// 좋은 예: Row-major coalesced 접근
__global__ void rowAccess(float* matrix, int M, int N) {
    int row = blockIdx.x * blockDim.x + threadIdx.x;
    float sum = 0.0f;
    for (int col = 0; col < N; col++) {
        sum += matrix[row * N + col];   // 연속 접근
    }
}

// 나쁜 예: 2-way bank conflict (stride=2)
__shared__ float smem[32][32];
// threadIdx.x=0 → smem[0][0], threadIdx.x=1 → smem[0][2]
// 짝수/홀수 스레드가 같은 bank에 충돌

// 좋은 예: 패딩으로 bank conflict 제거
__shared__ float smem_padded[32][33];  // 열에 1 패딩 추가
// threadIdx.x=0 → bank 0, threadIdx.x=1 → bank 1, ...

__global__ void matTransposeNoBankConflict(float* in, float* out, int N) {
    __shared__ float tile[TILE_SIZE][TILE_SIZE + 1]; // +1 패딩

    int x = blockIdx.x * TILE_SIZE + threadIdx.x;
    int y = blockIdx.y * TILE_SIZE + threadIdx.y;

    if (x < N && y < N)
        tile[threadIdx.y][threadIdx.x] = in[y * N + x];

    __syncthreads();

    x = blockIdx.y * TILE_SIZE + threadIdx.x;
    y = blockIdx.x * TILE_SIZE + threadIdx.y;

    if (x < N && y < N)
        out[y * N + x] = tile[threadIdx.x][threadIdx.y];
}

// Warp-level reduction with shuffle instructions
__device__ float warpReduceSum(float val) {
    for (int offset = 16; offset > 0; offset >>= 1) {
        val += __shfl_down_sync(0xffffffff, val, offset);
    }
    return val;
}

__global__ void blockReduceSum(float* input, float* output, int n) {
    __shared__ float shared[32];

    int tid   = threadIdx.x;
    int idx   = blockIdx.x * blockDim.x + tid;
    float val = (idx < n) ? input[idx] : 0.0f;

    // 워프 내 reduction
    val = warpReduceSum(val);

    // 각 워프의 결과를 shared memory에 저장
    if (tid % 32 == 0) shared[tid / 32] = val;
    __syncthreads();

    // 첫 번째 워프에서 최종 reduction
    if (tid < 32) {
        val = (tid < blockDim.x / 32) ? shared[tid] : 0.0f;
        val = warpReduceSum(val);
        if (tid == 0) atomicAdd(output, val);
    }
}

__global__ void exclusiveScan(int* input, int* output, int n) {
    __shared__ int temp[1024 * 2];

    int tid    = threadIdx.x;
    int offset = 1;

    // 데이터 로드
    temp[2 * tid]     = (2 * tid < n)     ? input[2 * tid]     : 0;
    temp[2 * tid + 1] = (2 * tid + 1 < n) ? input[2 * tid + 1] : 0;

    // Up-sweep (reduce) 단계
    for (int d = n >> 1; d > 0; d >>= 1) {
        __syncthreads();
        if (tid < d) {
            int ai = offset * (2 * tid + 1) - 1;
            int bi = offset * (2 * tid + 2) - 1;
            temp[bi] += temp[ai];
        }
        offset <<= 1;
    }

    // 루트를 0으로 초기화
    if (tid == 0) temp[n - 1] = 0;

    // Down-sweep 단계
    for (int d = 1; d < n; d <<= 1) {
        offset >>= 1;
        __syncthreads();
        if (tid < d) {
            int ai = offset * (2 * tid + 1) - 1;
            int bi = offset * (2 * tid + 2) - 1;
            int t  = temp[ai];
            temp[ai]  = temp[bi];
            temp[bi] += t;
        }
    }

    __syncthreads();
    if (2 * tid < n)     output[2 * tid]     = temp[2 * tid];
    if (2 * tid + 1 < n) output[2 * tid + 1] = temp[2 * tid + 1];
}

#include <mma.h>
using namespace nvcuda;

// WMMA 타일 크기 (Volta/Turing/Ampere 지원 크기)
#define WMMA_M 16
#define WMMA_N 16
#define WMMA_K 16

__global__ void tensorCoreMatMul(
    half* A, half* B, float* C,
    int M, int N, int K
) {
    // WMMA 프래그먼트 선언
    wmma::fragment<wmma::matrix_a, WMMA_M, WMMA_N, WMMA_K,
                   half, wmma::row_major> aFrag;
    wmma::fragment<wmma::matrix_b, WMMA_M, WMMA_N, WMMA_K,
                   half, wmma::col_major> bFrag;
    wmma::fragment<wmma::accumulator, WMMA_M, WMMA_N, WMMA_K,
                   float> cFrag;

    int warpM = (blockIdx.x * blockDim.x + threadIdx.x) / 32;
    int warpN = (blockIdx.y * blockDim.y + threadIdx.y);

    // accumulator 초기화
    wmma::fill_fragment(cFrag, 0.0f);

    // K 차원을 따라 타일 반복
    for (int i = 0; i < K; i += WMMA_K) {
        int aRow = warpM * WMMA_M;
        int aCol = i;
        int bRow = i;
        int bCol = warpN * WMMA_N;

        if (aRow < M && aCol < K && bRow < K && bCol < N) {
            // 행렬 프래그먼트 로드
            wmma::load_matrix_sync(aFrag, A + aRow * K + aCol, K);
            wmma::load_matrix_sync(bFrag, B + bRow * N + bCol, N);

            // Tensor Core MMA 실행
            wmma::mma_sync(cFrag, aFrag, bFrag, cFrag);
        }
    }

    // 결과 저장
    int cRow = warpM * WMMA_M;
    int cCol = warpN * WMMA_N;
    if (cRow < M && cCol < N) {
        wmma::store_matrix_sync(C + cRow * N + cCol, cFrag, N,
                                wmma::mem_row_major);
    }
}

__device__ float fastSqrt(float x) {
    float result;
    asm("sqrt.approx.f32 %0, %1;" : "=f"(result) : "f"(x));
    return result;
}

__device__ int ldgInt(const int* ptr) {
    int val;
    // 캐시 힌트를 사용한 글로벌 메모리 로드 (L2 캐시 우선)
    asm volatile("ld.global.cg.s32 %0, [%1];"
                 : "=r"(val) : "l"(ptr));
    return val;
}

표준 Attention:
  1. S = Q @ K^T        → HBM에 N×N 행렬 저장 (O(N²))
  2. P = softmax(S)     → HBM에서 읽고 쓰기
  3. O = P @ V          → HBM에서 읽고 쓰기

Flash Attention:
  - Q, K, V를 타일로 분할 → SRAM에 로드
  - 온라인 소프트맥스로 행 최댓값/분모를 점진적으로 갱신
  - HBM에는 최종 O만 저장 → HBM 읽기/쓰기 O(N) 달성

// Flash Attention 핵심 루프 (단순화된 버전)
__global__ void flashAttentionKernel(
    const float* Q, const float* K, const float* V,
    float* O, float* l, float* m,
    int seqLen, int headDim, float scale
) {
    int tileSize = blockDim.x;
    int qStart   = blockIdx.x * tileSize;

    extern __shared__ float smem[];
    float* sQ = smem;
    float* sK = smem + tileSize * headDim;
    float* sV = smem + 2 * tileSize * headDim;

    // Q 타일을 SRAM에 로드
    for (int d = 0; d < headDim; d++) {
        sQ[threadIdx.x * headDim + d] =
            Q[(qStart + threadIdx.x) * headDim + d];
    }

    float mi  = -INFINITY;  // 행 최댓값 (수치 안정성)
    float li  = 0.0f;       // 행 합산 (softmax 분모)
    float oi[64] = {};      // 출력 누적

    // K, V 타일을 순회
    for (int kStart = 0; kStart < seqLen; kStart += tileSize) {
        // K, V 타일 로드
        for (int d = 0; d < headDim; d++) {
            sK[threadIdx.x * headDim + d] =
                K[(kStart + threadIdx.x) * headDim + d];
            sV[threadIdx.x * headDim + d] =
                V[(kStart + threadIdx.x) * headDim + d];
        }
        __syncthreads();

        // S_ij = Q_i @ K_j^T * scale
        for (int j = 0; j < tileSize; j++) {
            float sij = 0.0f;
            for (int d = 0; d < headDim; d++) {
                sij += sQ[threadIdx.x * headDim + d]
                     * sK[j * headDim + d];
            }
            sij *= scale;

            // 온라인 softmax 갱신
            float miNew = fmaxf(mi, sij);
            float liNew = li * expf(mi - miNew) + expf(sij - miNew);

            for (int d = 0; d < headDim; d++) {
                oi[d] = oi[d] * (li / liNew) * expf(mi - miNew)
                      + (expf(sij - miNew) / liNew) * sV[j * headDim + d];
            }
            mi = miNew;
            li = liNew;
        }
        __syncthreads();
    }

    // 결과 저장
    for (int d = 0; d < headDim; d++) {
        O[(qStart + threadIdx.x) * headDim + d] = oi[d];
    }
    l[qStart + threadIdx.x] = li;
    m[qStart + threadIdx.x] = mi;
}

#include <cuda_runtime.h>

void pipelinedProcessing(float* h_input, float* h_output,
                         int totalSize, int chunkSize) {
    const int numStreams = 4;
    cudaStream_t streams[numStreams];
    float *d_input[numStreams], *d_output[numStreams];

    for (int i = 0; i < numStreams; i++) {
        cudaStreamCreate(&streams[i]);
        cudaMalloc(&d_input[i],  chunkSize * sizeof(float));
        cudaMalloc(&d_output[i], chunkSize * sizeof(float));
    }

    // 파이프라인: HtoD 복사 → 커널 → DtoH 복사가 중첩
    for (int chunk = 0; chunk < totalSize / chunkSize; chunk++) {
        int s = chunk % numStreams;
        float* hIn  = h_input  + chunk * chunkSize;
        float* hOut = h_output + chunk * chunkSize;

        // 비동기 Host-to-Device 복사
        cudaMemcpyAsync(d_input[s], hIn,
                        chunkSize * sizeof(float),
                        cudaMemcpyHostToDevice, streams[s]);

        // 커널 실행 (이전 memcpy와 다른 스트림은 병렬 수행)
        dim3 block(256);
        dim3 grid((chunkSize + 255) / 256);
        processKernel<<<grid, block, 0, streams[s]>>>(
            d_input[s], d_output[s], chunkSize);

        // 비동기 Device-to-Host 복사
        cudaMemcpyAsync(hOut, d_output[s],
                        chunkSize * sizeof(float),
                        cudaMemcpyDeviceToHost, streams[s]);
    }

    // 모든 스트림 완료 대기
    for (int i = 0; i < numStreams; i++) {
        cudaStreamSynchronize(streams[i]);
        cudaStreamDestroy(streams[i]);
        cudaFree(d_input[i]);
        cudaFree(d_output[i]);
    }
}

cudaEvent_t startEvent, stopEvent;
cudaEventCreate(&startEvent);
cudaEventCreate(&stopEvent);

cudaEventRecord(startEvent, stream1);
kernel1<<<grid, block, 0, stream1>>>(d_data);

// stream2가 stream1의 kernel1 완료를 기다림
cudaStreamWaitEvent(stream2, startEvent, 0);
kernel2<<<grid, block, 0, stream2>>>(d_data);

cudaEventRecord(stopEvent, stream2);
cudaEventSynchronize(stopEvent);

float milliseconds = 0;
cudaEventElapsedTime(&milliseconds, startEvent, stopEvent);

// GPU 0에서 GPU 1로 직접 메모리 복사 (NVLink/PCIe P2P)
cudaSetDevice(0);
float* d_src;
cudaMalloc(&d_src, size);

cudaSetDevice(1);
float* d_dst;
cudaMalloc(&d_dst, size);

// P2P 활성화 확인
int canAccess;
cudaDeviceCanAccessPeer(&canAccess, 0, 1);
if (canAccess) {
    cudaSetDevice(0);
    cudaDeviceEnablePeerAccess(1, 0);
    // GPU 0에서 GPU 1 메모리로 직접 복사
    cudaMemcpyPeer(d_dst, 1, d_src, 0, size);
}

#include <nccl.h>

void ncclAllReduceExample(int numGPUs, float** d_buffers, size_t count) {
    ncclComm_t comms[numGPUs];
    int devList[numGPUs];
    for (int i = 0; i < numGPUs; i++) devList[i] = i;

    // NCCL 통신자 초기화
    ncclCommInitAll(comms, numGPUs, devList);

    cudaStream_t streams[numGPUs];
    for (int i = 0; i < numGPUs; i++) {
        cudaSetDevice(i);
        cudaStreamCreate(&streams[i]);
    }

    // AllReduce: 모든 GPU의 버퍼를 합산하고 결과를 모든 GPU에 분배
    ncclGroupStart();
    for (int i = 0; i < numGPUs; i++) {
        cudaSetDevice(i);
        ncclAllReduce(d_buffers[i], d_buffers[i], count,
                      ncclFloat, ncclSum, comms[i], streams[i]);
    }
    ncclGroupEnd();

    for (int i = 0; i < numGPUs; i++) {
        cudaSetDevice(i);
        cudaStreamSynchronize(streams[i]);
        ncclCommDestroy(comms[i]);
    }
}

import triton
import triton.language as tl
import torch

@triton.jit
def matmul_kernel(
    a_ptr, b_ptr, c_ptr,
    M, N, K,
    stride_am, stride_ak,
    stride_bk, stride_bn,
    stride_cm, stride_cn,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

    # 오프셋 계산
    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    offs_k = tl.arange(0, BLOCK_K)

    # 포인터 초기화
    a_ptrs = a_ptr + (offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak)
    b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn)

    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)

    for k in range(0, K, BLOCK_K):
        # 경계 마스킹으로 안전하게 로드
        a = tl.load(a_ptrs, mask=(offs_m[:, None] < M) & (offs_k[None, :] < K - k), other=0.0)
        b = tl.load(b_ptrs, mask=(offs_k[:, None] < K - k) & (offs_n[None, :] < N), other=0.0)

        acc += tl.dot(a, b)

        a_ptrs += BLOCK_K * stride_ak
        b_ptrs += BLOCK_K * stride_bk

    # 결과 저장
    c_ptrs = c_ptr + stride_cm * offs_m[:, None] + stride_cn * offs_n[None, :]
    c_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
    tl.store(c_ptrs, acc, mask=c_mask)


def matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    M, K = a.shape
    K2, N = b.shape
    assert K == K2

    c = torch.empty((M, N), device=a.device, dtype=torch.float32)

    grid = lambda meta: (
        triton.cdiv(M, meta['BLOCK_M']),
        triton.cdiv(N, meta['BLOCK_N']),
    )

    matmul_kernel[grid](
        a, b, c,
        M, N, K,
        a.stride(0), a.stride(1),
        b.stride(0), b.stride(1),
        c.stride(0), c.stride(1),
        BLOCK_M=128, BLOCK_N=256, BLOCK_K=64,
    )
    return c

@triton.jit
def flash_attn_fwd_kernel(
    Q_ptr, K_ptr, V_ptr, O_ptr,
    L_ptr, M_ptr,
    seqlen, headdim, scale,
    BLOCK_Q: tl.constexpr,
    BLOCK_KV: tl.constexpr,
):
    start_q = tl.program_id(0) * BLOCK_Q
    offs_q  = start_q + tl.arange(0, BLOCK_Q)
    offs_d  = tl.arange(0, headdim)

    # Q 블록 로드
    q = tl.load(Q_ptr + offs_q[:, None] * headdim + offs_d[None, :],
                mask=(offs_q[:, None] < seqlen), other=0.0)

    mi  = tl.full((BLOCK_Q,), value=-float('inf'), dtype=tl.float32)
    li  = tl.zeros((BLOCK_Q,), dtype=tl.float32)
    acc = tl.zeros((BLOCK_Q, headdim), dtype=tl.float32)

    for start_kv in range(0, seqlen, BLOCK_KV):
        offs_kv = start_kv + tl.arange(0, BLOCK_KV)

        k = tl.load(K_ptr + offs_kv[:, None] * headdim + offs_d[None, :],
                    mask=(offs_kv[:, None] < seqlen), other=0.0)
        v = tl.load(V_ptr + offs_kv[:, None] * headdim + offs_d[None, :],
                    mask=(offs_kv[:, None] < seqlen), other=0.0)

        # Attention score 계산
        s = tl.dot(q, tl.trans(k)) * scale

        # 온라인 softmax
        mi_new = tl.maximum(mi, tl.max(s, axis=1))
        p      = tl.exp(s - mi_new[:, None])
        li_new = li * tl.exp(mi - mi_new) + tl.sum(p, axis=1)

        acc = acc * (li * tl.exp(mi - mi_new))[:, None] / li_new[:, None]
        acc += tl.dot(p, v) / li_new[:, None]

        mi = mi_new
        li = li_new

    # 결과 저장
    tl.store(O_ptr + offs_q[:, None] * headdim + offs_d[None, :],
             acc, mask=(offs_q[:, None] < seqlen))
    tl.store(L_ptr + offs_q, li, mask=(offs_q < seqlen))
    tl.store(M_ptr + offs_q, mi, mask=(offs_q < seqlen))

# 커널 상세 분석
ncu --set full -o profile_report ./my_cuda_app

# 특정 커널만 프로파일링
ncu --kernel-name tiledMatMul --launch-count 3 ./my_cuda_app

# 메모리 대역폭 분석 (Roofline 모델용)
ncu --metrics \
    l1tex__t_bytes_pipe_lsu_mem_global_op_ld.sum,\
    l1tex__t_bytes_pipe_lsu_mem_global_op_st.sum,\
    sm__cycles_elapsed.avg,\
    sm__cycles_elapsed.avg.per_second \
    ./my_cuda_app

# DRAM 대역폭 활용률
ncu --metrics \
    gpu__dram_throughput.avg.pct_of_peak_sustained_elapsed \
    ./my_cuda_app

# 전체 애플리케이션 타임라인 수집
nsys profile --trace=cuda,nvtx,osrt \
             --output=timeline_report \
             ./my_cuda_app

# 통계 요약 출력
nsys stats timeline_report.nsys-rep

# GPU 활동도 분석
nsys analyze --report gputrace timeline_report.nsys-rep

#include <nvtx3/nvToolsExt.h>

void profiledFunction() {
    nvtxRangePush("DataPreprocessing");
    preprocessData();
    nvtxRangePop();

    nvtxRangePush("ForwardPass");
    forwardPass();
    nvtxRangePop();
}

산술 강도(Arithmetic Intensity) = FLOP / Bytes

예시 (A100 80GB):
- 피크 FP16 Tensor Core: 312 TFLOPS
- HBM 대역폭: 2 TB/s
- Ridge Point: 312 / 2 = 156 FLOP/Byte

커널의 산술 강도 < 156 → 메모리 바운드
커널의 산술 강도 > 156 → 연산 바운드

# 커널 산술 강도 측정
ncu --metrics \
    sm__ops_path_tensor_src_fp16.sum,\
    dram__bytes.sum \
    ./matmul_app
# 산술 강도 = tensor_ops * 2 / dram_bytes