Benchmark Script, 70B LLM Hardware Sizing மற்றும் AI Compiler Stacks

நவீன AI அமைப்புப் பொறியியல் (systems engineering), kernel-மட்ட benchmarking, கொள்ளளவு மற்றும் bandwidth கணிதம், மற்றும் பல-கட்ட compiler தாழ்த்தல் (lowering) ஆகியவற்றை ஒன்றிணைக்கிறது. கீழே மூன்று துறைகளினதும் முழுமையான பகுப்பாய்வு உள்ளது.

1. முழுமையான Benchmark Script: Eager PyTorch எதிர் torch.compile எதிர் Custom Triton

GPU இயக்கம் ஒத்திசைவற்றது (asynchronous) என்பதால், துல்லியமான CUDA benchmarking-க்கு warmup pass-கள் (JIT compilation-ஐத் தூண்டி memory cache-களை ஒதுக்குவதற்கு) மற்றும் wall-clock timer-களை வாசிப்பதற்கு முன் torch.cuda.synchronize() (அல்லது CUDA Events) தேவை.

கீழே 4096 × 8192 FP16 matrix-இல் fused row-wise Softmax பணிச்சுமையில் Eager PyTorch, torch.compile (TorchInductor), மற்றும் Custom Triton Kernel ஆகியவற்றை ஒப்பிடும் தனித்த (standalone) Python script உள்ளது.

import torch
import triton
import triton.language as tl

# 1. Custom Fused Triton Kernel
@triton.jit
def fused_softmax_kernel(
    output_ptr, input_ptr,
    input_row_stride, output_row_stride,
    n_cols,
    BLOCK_SIZE: tl.constexpr
):
    row_idx = tl.program_id(0)
    row_start_ptr = input_ptr + row_idx * input_row_stride
    col_offsets = tl.arange(0, BLOCK_SIZE)
    mask = col_offsets < n_cols

    # Single read from HBM into SRAM/registers
    row = tl.load(row_start_ptr + col_offsets, mask=mask, other=-float('inf')).to(tl.float32)
    row_minus_max = row - tl.max(row, axis=0)
    numerator = tl.exp(row_minus_max)
    denominator = tl.sum(numerator, axis=0)
    softmax_output = (numerator / denominator).to(tl.float16)

    # Single write back to HBM
    output_row_start_ptr = output_ptr + row_idx * output_row_stride
    tl.store(output_row_start_ptr + col_offsets, softmax_output, mask=mask)

def triton_softmax(x: torch.Tensor) -> torch.Tensor:
    n_rows, n_cols = x.shape
    BLOCK_SIZE = triton.next_power_of_2(n_cols)
    y = torch.empty_like(x)
    fused_softmax_kernel[(n_rows,)](
        y, x, x.stride(0), y.stride(0), n_cols, BLOCK_SIZE=BLOCK_SIZE
    )
    return y

# 2. Unfused Eager PyTorch Baseline (Forces multiple HBM round-trips)
def naive_pytorch_softmax(x: torch.Tensor) -> torch.Tensor:
    x_fp32 = x.float()
    x_max = x_fp32.max(dim=-1, keepdim=True)[0]
    exp_x = torch.exp(x_fp32 - x_max)
    return (exp_x / exp_x.sum(dim=-1, keepdim=True)).half()

# 3. TorchInductor Compiled Function
compiled_softmax = torch.compile(naive_pytorch_softmax, mode="max-autotune")

# 4. Accurate CUDA Event Benchmarking Harness
def benchmark_fn(fn, x, warmup=25, rep=100):
    for _ in range(warmup):
        fn(x)
    torch.cuda.synchronize()

    start_events = [torch.cuda.Event(enable_timing=True) for _ in range(rep)]
    end_events = [torch.cuda.Event(enable_timing=True) for _ in range(rep)]

    for i in range(rep):
        start_events[i].record()
        fn(x)
        end_events[i].record()

    torch.cuda.synchronize()
    times_ms = [s.elapsed_time(e) for s, e in zip(start_events, end_events)]
    median_ms = sorted(times_ms)[rep // 2]

    # Effective Memory Bandwidth (GB/s): 1 read + 1 write of FP16 tensor
    bytes_transferred = 2 * x.numel() * x.element_size()
    gb_per_sec = (bytes_transferred / (median_ms * 1e-3)) / 1e9
    return median_ms, gb_per_sec

if __name__ == "__main__":
    assert torch.cuda.is_available(), "CUDA GPU required."
    x = torch.randn((4096, 8192), device="cuda", dtype=torch.float16)

    # Verify numerical accuracy against native PyTorch
    ref = torch.softmax(x, dim=-1)
    assert torch.allclose(triton_softmax(x), ref, atol=1e-2)
    assert torch.allclose(compiled_softmax(x), ref, atol=1e-2)

    for name, fn in [
        ("Naive Eager PyTorch", naive_pytorch_softmax),
        ("torch.compile (Inductor)", compiled_softmax),
        ("Custom Triton Kernel", triton_softmax),
    ]:
        ms, bw = benchmark_fn(fn, x)
        print(f"{name:<26} | Latency: {ms:.4f} ms | Effective BW: {bw:.1f} GB/s")

• முடிவுகள் காட்டுவது: H100-இல் (3,350 GB/s உச்ச HBM bandwidth), Naive Eager PyTorch வழக்கமாக ~600–800 GB/s பயனுள்ள bandwidth-ஐ மட்டுமே அடைகிறது, ஏனெனில் அது 5 தனித்தனி kernel-களை launch செய்து இடைநிலை tensor-களை VRAM-இல் திரும்பத்திரும்ப எழுதுகிறது. torch.compile-ம் Custom Triton Kernel-ம் முழுச் சங்கிலியையும் ஒரே pass-ஆக இணைத்து, வழக்கமாக 2,600–3,000+ GB/s-ஐ அடைகின்றன (கோட்பாட்டு hardware உச்சத்தின் 80–90%).

2. Hardware Sizing திட்டம்: Production-இல் 70B LLM-ஐ வழங்குதல்

Cluster-ஐத் துல்லியமாக அளவிட (size), நிலையான weight நினைவகம், இயங்கும் (dynamic) KV cache நினைவகம், மற்றும் memory-bandwidth-bound decode throughput ஆகியவற்றைக் கணக்கிட வேண்டும்.

இலக்கு பணிச்சுமை அனுமானங்கள் (Target Workload Assumptions)

• Model: Llama-3-பாணி 70B Dense Model (80 layers, GQA வழியே 8 KV heads, 128 head dimension).

• Precision: FP8 weights (param ஒன்றுக்கு 1 byte) மற்றும் FP8 KV Cache (உறுப்பொன்றுக்கு (element) 1 byte).

• ஒரேசமைப் பயனாளர்கள் & Context: 64 ஒரேசமைப் (concurrent) பயனாளர்கள், ஒவ்வொருவரும் 8,192-token context window உடன்.

• Hardware வாய்ப்பாளர்: NVIDIA H100 80GB SXM (80 GB HBM3, 3.35 TB/s memory bandwidth).

படி 1: நிலையான Model Weight VRAM-ஐக் கணக்கிடுதல்

Weight VRAM=70× 10^(9) params× 1 byte (FP8)≈ 70 GB

படி 2: இயங்கும் KV Cache VRAM-ஐக் கணக்கிடுதல்

ஒவ்வொரு token-க்கும், ஒவ்வொரு layer-ம் 128 பரிமாண (dimension) 8 GQA head-களில் ஒரு Key vector-ஐயும் ஒரு Value vector-ஐயும் சேமிக்கிறது:

KV Bytes per Token = 2 (K & V) × 80 layers × 8 KV heads × 128 dim × 1 byte = 163,840 bytes ≈ 160 KB/token

ஒவ்வொருவரும் 8,192 token-களுடன் 64 ஒரேசமைப் பயனாளர்களுக்கு (மொத்தம் 524,288 செயலில் உள்ள token-கள்):

Total KV Cache VRAM=524,288 tokens× 163,840 bytes≈ 85.9 GB

படி 3: GPU எண்ணிக்கை & Tensor Parallelism-ஐத் தீர்மானித்தல்

Total Required VRAM=70 GB (Weights)+85.9 GB (KV Cache)+~15 GB (CUDA Graphs/Activations)≈ 171 GB

• குறைந்தபட்ச Hardware: NVLink வழியே இணைக்கப்பட்ட 4× H100 80GB GPU-கள் (மொத்தம் 320 GB HBM), Tensor Parallelism = 4 (TP=4) உடன் இயங்கும்.

• GPU ஒன்றுக்கான Memory Footprint: ஒவ்வொரு GPU-யும் 17.5 GB பிரிக்கப்பட்ட (sharded) weights + 21.5 GB KV cache + ~4 GB runtime மேல்சுமை = 80 GB-இல் ~43 GB பயன்படுத்தப்பட்டு, போக்குவரத்து உச்சங்களுக்கு (traffic spikes) நிறைய இடவசதியை (headroom) விட்டுச் செல்கிறது.

படி 4: Token உருவாக்க Throughput-ஐ மதிப்பிடுதல் (Decode வேகம்)

Decode-இன் போது, 64 பயனாளர்களுக்கும் ஒரு படியை உருவாக்க, முழு 70 GB model weights-ஐயும் 85.9 GB KV cache-ஐயும் ஒருமுறை HBM-இலிருந்து வாசிக்க வேண்டும் (ஒரு படிக்கு மொத்தம் 155.9 GB வாசிப்பு), 4 GPU-களின் இணைந்த memory bandwidth முன்னிலையில் (4 × 3.35 TB/s = 13.4 TB/s, ~75% bus திறன் எனக் கருதினால் = 10.05 TB/s பயனுள்ள):

Time per Decode Step (TPOT)≈ (155.9 GB)/(10,050 GB/s)≈ 15.5 ms per token

Per-User Speed≈ (1000 ms)/(15.5 ms)≈ 64.5 tokens/sec/user | Total System Throughput=64× 64.5≈ 4,128 tokens/sec

3. AI Compiler Stacks: Python-ஐ PTX Assembly வரை தாழ்த்தல் (Lowering)

நவீன AI compiler-கள், பல-கட்ட Intermediate Representation (IR) pipeline வழியே உயர்-மட்ட இயங்கு (dynamic) Python code-ஐ bare-metal GPU இயந்திர அறிவுறுத்தல்களாக மாற்றுகின்றன.

• MLIR (Multi-Level Intermediate Representation) நவீன AI Compiler-களில் ஏன் ஆதிக்கம் செலுத்துகிறது.

• வரலாற்றில், நிலையான LLVM போன்ற compiler-கள் தாழ்-மட்ட scalar loop-களிலும் pointer-களிலும் மட்டுமே இயங்கின; உயர்-மட்ட matrix அமைப்பைக் கண்முன் இழந்துவிட்டன.
• MLIR, compiler-களுக்கு (Triton, XLA, Mojo போன்றவை) ஒரேசமயம் பல மட்ட சுருக்கங்களை (abstraction) வைத்திருக்க அனுமதிக்கிறது — LLVM IR, PTX (Parallel Thread Execution) virtual assembly, இறுதியாக SASS (உங்கள் குறிப்பிட்ட NVIDIA GPU தலைமுறை இயக்கும் கச்சிதமான binary இயந்திரக் குறியீடு) வரை படிப்படியாகத் தாழ்த்துவதற்கு முன், உயர் dialect-களில் 2D matrix tile-களையும் shared memory அமைப்புகளையும் (layouts) பற்றிச் சிந்திக்கிறது.