ජනනය කළ Triton කේතය, GEMM උදාහරණ සහ ROCm warp-size අන්තරාය

මෙම ප්‍රායෝගික යාන්ත්‍රණ තුන—torch.compile ගේ ජනනය කළ Triton කේතය පරීක්ෂා කිරීම, tiled matrix multiplication (GEMM) කේත සංසන්දනය සහ CUDA AMD ROCm වෙත ගෙන යන විට warp-size අන්තරායන් හැසිරවීම—ඉහළ-මට්ටමේ AI frameworks සෘජුවම දෘඩාංග උපදෙස් බවට පරිවර්තනය වන ආකාරය නිදර්ශනය කරයි.

1. torch.compile විසින් ජනනය කළ Triton කේතය පරීක්ෂා කිරීම

ඔබ RMSNorm (x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) * weight) වැනි සම්මත PyTorch මෙහෙයුමක් @torch.compile සමඟ ආවරණය කරන විට, TorchInductor පසුතලයෙන් තනි fused Triton kernel එකක් ජනනය කිරීමෙන් වෙනම HBM round-trips හතරක් ඉවත් කරයි.

ඔබේ script එක TORCH_LOGS="output_code" python script.py සමඟ ධාවනය කිරීමෙන් සැබෑ compiler-ජනිත Triton කේතය මුද්‍රණය වේ, එය මෙසේ පෙනේ:

import triton
import triton.language as tl

@triton.jit
def triton_per_fused_mul_pow_rsqrt_0(in_ptr0, in_ptr1, out_ptr0, xnumel, rnumel, XBLOCK: tl.constexpr):
    RBLOCK: tl.constexpr = 4096
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
    xmask = xindex < xnumel
    rindex = tl.arange(0, RBLOCK)[None, :]
    rmask = rindex < rnumel

    # 1. Single coalesced read of input tensor (x) and weight vector from HBM
    x = tl.load(in_ptr0 + (rindex + (4096 * xindex)), rmask & xmask, other=0.0).to(tl.float32)
    w = tl.load(in_ptr1 + (rindex), rmask, eviction_policy='evict_last').to(tl.float32)

    # 2. Compute square, row-wise mean, and reciprocal square root entirely in registers
    x_sq = x * x
    row_sum = tl.sum(x_sq, 1)[:, None]
    rstd = tl.rsqrt((row_sum / 4096.0) + 1e-06)

    # 3. Multiply by inverse standard deviation and learned weight, then write once
    out = (x * rstd) * w
    tl.store(out_ptr0 + (rindex + (4096 * xindex)), out, rmask & xmask)

• Compiler එක ස්වයංක්‍රීයව සොයා ගත් දේ: weight vector w මත eviction_policy='evict_last' සැලකිල්ලට ගන්න. TorchInductor w සෑම row එකකම (xindex) නැවත භාවිතා වන බව ස්වයංක්‍රීයව හඳුනා ගෙන, x එක වරක් ගලා යන අතරතුර w on-chip තබා ගැනීමට GPU ගේ L2 cache එකට උපදෙස් දුන්නේය.

2. පැති-පැති Matrix Multiplication (GEMM) කේතය: CUDA C++ එදිරිව Triton

Matrix multiplication (C = A × B) CUDA C++ හි අතින් thread synchronization සහ Triton හි ප්‍රකාශනීය tile ගණිතය අතර ඇති පැහැදිලි වෙනස ඉස්මතු කරයි.

A. CUDA C++ (පැහැදිලි Shared Memory Tiling & Sync Barriers)

CUDA C++ තුළ, සෑම thread එකක්ම එක් scalar ප්‍රතිදාන මූලද්‍රව්‍යයක් (acc) ගණනය කරන අතර, threads __shared__ memory වෙත tiles පූරණය කිරීමට අතින් සහයෝගයෙන් කටයුතු කරයි:

#define TILE 16

__global__ void gemm_cuda(const float* A, const float* B, float* C, int M, int N, int K) {
    __shared__ float sA[TILE][TILE];
    __shared__ float sB[TILE][TILE];

    int row = blockIdx.y * TILE + threadIdx.y;
    int col = blockIdx.x * TILE + threadIdx.x;
    float acc = 0.0f;

    for (int t = 0; t < (K + TILE - 1) / TILE; ++t) {
        // 1. Each thread manually loads 1 element into Shared Memory
        sA[threadIdx.y][threadIdx.x] = (row < M && t * TILE + threadIdx.x < K)
                                       ? A[row * K + t * TILE + threadIdx.x] : 0.0f;
        sB[threadIdx.y][threadIdx.x] = (t * TILE + threadIdx.y < K && col < N)
                                       ? B[(t * TILE + threadIdx.y) * N + col] : 0.0f;
        __syncthreads(); // Mandatory barrier before reading shared tiles

        // 2. Inner dot product across the 16x16 tile
        for (int k = 0; k < TILE; ++k) {
            acc += sA[threadIdx.y][k] * sB[k][threadIdx.x];
        }
        __syncthreads(); // Mandatory barrier before loading next tile
    }

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

B. OpenAI Triton (ප්‍රකාශනීය 2D Tile Accumulation & Tensor Cores)

Triton තුළ, තනි program block එකක් tl.dot() භාවිතයෙන් සම්පූර්ණ BLOCK_M × BLOCK_N ප්‍රතිදාන tile එකක් ගණනය කරන අතර, එය ස්වයංක්‍රීයව දෘඩාංග Tensor Cores දක්වා පහත හෙළයි:

@triton.jit
def gemm_triton(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)

    # Accumulate in FP32 registers while multiplying FP16 tiles on Tensor Cores
    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) & (k + offs_k[None, :] < K), other=0.0)
        b = tl.load(b_ptrs, mask=(k + offs_k[:, None] < K) & (offs_n[None, :] < N), other=0.0)
        acc = tl.dot(a, b, acc)  # Hardware Tensor Core instruction
        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, :]
    tl.store(c_ptrs, acc.to(tl.float16), mask=(offs_m[:, None] < M) & (offs_n[None, :] < N))

3. CUDA C++ AMD ROCm (HIP) වෙත ගෙන යන විට Warp-Size අන්තරායන්

පහළ-මට්ටමේ CUDA C++ kernels AMD ගේ HIP interface එක හරහා AMD Instinct GPUs (MI300X හෝ MI325X වැනි) වෙත ගෙන යන විට, syntax පරිවර්තනය (cudaMalloc → hipMalloc) සුළු වේ. කෙසේ වෙතත්, NVIDIA warps වල threads 32 ක් අඩංගු වන අතර AMD CDNA wavefronts වල threads 64 ක් අඩංගු බැවින් නිහඬ තාර්කික දෝෂ නිතර ඇති වේ.

• Triton අන්තරායන් හතරම වළක්වන හේතුව: OpenAI Triton හි, tl.sum() සහ tl.max() වැනි මෙහෙයුම් hardcoded 32-thread warps වලට වඩා සම්පූර්ණ තාර්කික tensor මාන හරහා ක්‍රියාත්මක වේ. සංයුක්තකරණය අතරතුර, Triton ගේ NVPTX backend එක NVIDIA GPUs සඳහා ස්වයංක්‍රීයව 32-thread shuffle trees නිකුත් කරන අතර, එහි AMDGPU backend එක AMD Instinct GPUs සඳහා 64-thread wavefront reductions නිකුත් කරයි—මෙය warp-size දෝෂ සැලසුම් අනුවම ඉවත් කරයි.

අපි දැන් නවීන GPU kernels හි සම්පූර්ණ ජීවන චක්‍රය—PyTorch bytecode ග්‍රහණය සහ Triton කේත ජනනය සිට CUDA එදිරිව AMD wavefront යාන්ත්‍රණය දක්වා—පරීක්ෂා කළෙමු.