LLM අනුමානය, ඒකාබද්ධ කර්නල් සහ GPU ජාලකරණය

නවීන මහා පරිමාණ AI අනුමානය මිලි තත්පරවලින් ටෝකන ජනනය කිරීමට සම්බන්ධීකරණය කළ බහු-GPU ක්‍රියාත්මක කිරීමේ නළමාර්ග, SRAM තුළ ඒකාබද්ධ කළ Triton/CUDA කර්නල් සහ ඉතා අඩු ප්‍රමාද සහිත fabric අන්තර්සම්බන්ධතා එකට යොදා ගනී.

1. මුල සිට අග දක්වා අනුමාන ක්‍රියාවලිය: ප්‍රොම්ප්ට් සිට ටෝකනය දක්වා

පරිශීලකයකු, vLLM වැනි අනුමාන එන්ජිමක් තුළ Tensor Parallelism යොදා ගනිමින් GPU 4ක් පුරා ධාවනය වන 70B-පරාමිති LLM එකකට ටෝකන 1,000ක ප්‍රොම්ප්ට් එකක් යවන බව සිතන්න.

• පියවර 1: ටෝකනීකරණය සහ සත්කාරකයේ සිට උපාංගයට මාරු කිරීම (CPU → GPU).

• CPU ටෝකනීකාරකය පෙළ පූර්ණ සංඛ්‍යා ටෝකන ID 1,000ක අරාවක් බවට පරිවර්තනය කරයි.
• සේවා එන්ජිම PagedAttention හරහා GPU HBM තුළ අඛණ්ඩ නොවන භෞතික මතක බ්ලොක් වෙන් කර, PCIe බසය හරහා ටෝකන ID GPU 0 හි VRAM වෙත පිටපත් කරයි.

• පියවර 2: Embedding සොයා බැලීම සහ විකාශනය (NVLink).

• GPU 0, embedding දෛශික 1,000 සොයා බලයි (උදා., හැඩය [1000, 8192]) සහ GPU 4ටම සමාන ආරම්භක සක්‍රියතා තිබෙන පරිදි NVLink හරහා ටෙන්සරය විකාශනය කරයි.

• පියවර 3: පෙර-පිරවුම් අදියර (ගණනයෙන් සීමා වන Tensor Core ක්‍රියාත්මක කිරීම).

• Transformer ස්තර 80න් එක් එක් ස්තරය තුළ, FlashAttention ප්‍රොම්ප්ට් ටෝකන 1,000ම හරහා එකවර ස්වයං-අවධානය ගණනය කරයි. ලැබෙන Key (K) සහ Value (V) ටෙන්සර් එක් එක් GPU හි PagedAttention KV හැඹිලි බ්ලොක්වල ලියනු ලැබේ.
• Feed-Forward Network (FFN) බර GPU 4 පුරා තීරු අනුව බෙදා තිබෙන නිසා (Tensor Parallelism), එක් එක් GPU එක FP8/FP4 Tensor Cores භාවිත කරමින් [1000, 8192] සක්‍රියතා තම දේශීය බර කොටසෙන් ගුණ කරයි.
• ස්තරය අවසානයේ GPU 4 තම අර්ධ ප්‍රතිඵල එකතු කිරීමට NVLink All-Reduce සිදු කරයි. එමඟින් ඊළඟ ස්තරයට ඇතුළු වීමට පෙර සෑම GPU එකකම සම්පූර්ණ යාවත්කාලීන කළ [1000, 8192] ටෙන්සරය පවතී.

• පියවර 4: නියැදීම සහ පළමු ටෝකන ප්‍රතිදානය.

• ස්තරය 80න් පසු, ටෝකන 1,000 හි අවසන් සැඟවුණු තත්ත්වය වචන මාලා න්‍යාසයෙන් ගුණ කර, වචන මාලාවේ වචන ~128,000ක් සඳහා logits සාදයි. නියැදි කර්නලයක් පළමු ප්‍රතිදාන ටෝකනය තෝරා එය CPU ප්‍රවාහයට ආපසු යවයි.

• පියවර 5: ස්වයංප්‍රතිගාමී විකේතන ලූපය (CUDA Graphs + මතකයෙන් සීමා වන ක්‍රියාත්මක කිරීම).

• ටෝකන 1,001 ජනනය කිරීමට CPU එක තනි තනි කර්නල් සිය ගණනක් නැවත දියත් නොකරයි. ඒ වෙනුවට තනි API ඇමතුමකින් පෙර පටිගත කළ CUDA Graph එකක් ක්‍රියාත්මක කරයි.
• දැන් තනි [1, 8192] ටෝකන දෛශිකයක් පමණක් සකසන සෑම GPU එකක්ම, තම දේශීය මාදිලි බර සහ හැඹිලිගත KV දෛශික 1,000ම HBM සිට චිපය-මත SRAM වෙත ප්‍රවාහනය කරයි, නව අවධාන ප්‍රතිදානය ගණනය කරයි, NVLink All-Reduce ක්‍රියාත්මක කරයි, සහ ටෝකන 1,001 හි K සහ V දෛශික ඊළඟ හිස් PagedAttention ස්ථානයට එක් කරයි. End-of-Sequence (EOS) ටෝකනයක් ජනනය වන තෙක් මෙම ලූපය නැවත නැවත ක්‍රියාත්මක වේ.

2. OpenAI Triton තුළ අභිරුචි ඒකාබද්ධ Softmax කර්නලයක් ගොඩනැඟීම

සරල PyTorch Softmax (torch.exp(x) / torch.sum(torch.exp(x))) සඳහා වෙන්වූ HBM මතක වාර තුනක් අවශ්‍ය වේ (සංඛ්‍යාත්මක ස්ථායීතාව සඳහා පේළියේ උපරිමය සොයන එකක්, ඝාතීය අගයන් ගණනය කර එකතු කරන එකක් සහ බෙදීම සඳහා එකක්).

පහත දැක්වෙන්නේ OpenAI Triton හි සම්පූර්ණ, නිෂ්පාදන-ශෛලියේ Fused Row-Wise Softmax කර්නලයකි. එය සම්පූර්ණ පේළියක් එක් වරක් වේගවත් චිපය-මත SRAM තුළට පූරණය කර, GPU රෙජිස්ටර ඇතුළත පියවර තුනම සිදු කර, එක් මතක වාරයකින් අවසන් ප්‍රතිදානය HBM වෙත නැවත ලියයි.

import torch
import triton
import triton.language as tl

@triton.jit
def fused_softmax_kernel(
    output_ptr, input_ptr,
    input_row_stride, output_row_stride,
    n_cols,
    BLOCK_SIZE: tl.constexpr
):
    # Each program instance processes one row of the matrix
    row_idx = tl.program_id(0)
    
    # Calculate memory pointers for the start of this specific row
    row_start_ptr = input_ptr + row_idx * input_row_stride
    col_offsets = tl.arange(0, BLOCK_SIZE)
    input_ptrs = row_start_ptr + col_offsets
    
    # Mask out-of-bounds threads if n_cols < BLOCK_SIZE (padding with -inf)
    mask = col_offsets < n_cols
    row = tl.load(input_ptrs, mask=mask, other=-float('inf'))
    
    # Step 1: Subtract row max inside fast SRAM/registers for numerical stability
    row_minus_max = row - tl.max(row, axis=0)
    
    # Step 2: Compute exponentials and sum in registers (zero HBM traffic)
    numerator = tl.exp(row_minus_max)
    denominator = tl.sum(numerator, axis=0)
    
    # Step 3: Normalize and write final probabilities back to HBM once
    softmax_output = numerator / denominator
    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
    # Round block size up to the next power of 2 to fit the entire row in SRAM
    BLOCK_SIZE = triton.next_power_of_2(n_cols)
    y = torch.empty_like(x)
    
    # Launch 1 Triton program block per matrix row
    fused_softmax_kernel[(n_rows,)](
        y, x,
        x.stride(0), y.stride(0),
        n_cols,
        BLOCK_SIZE=BLOCK_SIZE
    )
    return y

3. AI පොකුරු ජාලකරණය: NVLink, InfiniBand සහ RoCEv2

AI වැඩබර GPU දහස් ගණනක් පුරා පරිමාණනය කරන විට, ජාල ප්‍රමාදය සහ තදබදය අමු GPU FLOPS තරමටම වැදගත් වේ. නවීන පොකුරු ද්වි-ස්තර scale-up සහ scale-out ජාල ධුරාවලියක් භාවිත කරයි.

• සාමූහික මෙහෙයුම් ජාලය භාවිත කරන ආකාරය.

• All-Reduce (අනුක්‍රමණ සහ ටෙන්සර් සමමුහුර්තකරණය): සෑම GPU එකක්ම අර්ධ ටෙන්සරයකින් ආරම්භ කර, පොකුර පුරා සියලු ටෙන්සර් එකතු කර, අවසානයේ සමාන අවසන් එකතුවක් ලබා ගනී. Ring හෝ Tree ඇල්ගොරිතම, නැතහොත් ස්විච් තම සිලිකන් තුළම දෛශික සෘජුව එක් කරන In-Network Computing (SHARP) භාවිත කළ විට, කලාප පළල පොකුරු ප්‍රමාණයෙන් ස්වාධීනව පරිමාණනය වේ.
• All-to-All (Mixture-of-Experts මාර්ගගත කිරීම): සෑම GPU එකක්ම අසල්වැසියන්ට එකම දත්ත යවන All-Reduce මෙන් නොව, All-to-All තුළ සෑම GPU එකක්ම ඉලක්කගත experts රඳවා ඇති අනෙක් සෑම GPU එකකටම අනන්‍ය ටෝකන කොටස් එකවර යවමින් සියල්ලෙන්-සියල්ලට දත්ත ප්‍රවාහයේ හදිසි ඉහළ යාමක් ඇති කරයි. මෙය ලබන්නාගේ වරායවල දැඩි incast තදබදයක් ඇති කරයි.

• AI ජාලවලට අලාභ රහිත RDMA අවශ්‍ය ඇයි.

• සාම්ප්‍රදායික TCP/IP ජාලකරණය පැකට් සැකසීමට CPU මත රඳා පවතින අතර බෆර පිරුණු විට පැකට් අත්හරී. එයින් මිලි තත්පර මට්ටමේ නැවත යැවීම් ඇති වී GPU 10,000ක සම්පූර්ණ පුහුණු පොකුරම නතර වේ.
• InfiniBand සහ RoCEv2 යන දෙකම RDMA (Remote Direct Memory Access) භාවිත කරයි. එමඟින් එක් GPU එකක Network Interface Card (NIC) හට CPU මැදිහත්වීමක් කිසිසේත් නොමැතිව වෙනත් GPU එකක HBM VRAM තුළට සෘජුව කියවීමට/ලිවීමට ඉඩ ලැබේ. RoCEv2 සම්මත Ethernet මත ශූන්‍ය පැකට් අලාභයක් බලාත්මක කරන්නේ, ස්විච් බෆර උතුරා යාමට පෙර ඉහළ ප්‍රවාහයේ යවන්නන් නතර කිරීමට PFC (Priority Flow Control) සහ සම්ප්‍රේෂණ වේග සුමටව සීමා කිරීමට ECN (Explicit Congestion Notification) යොදා ගනිමිනි.