End-to-End Inference Walkthrough, Triton Fused Softmax Kernel மற்றும் AI Cluster Networking

நவீன பெரிய-அளவு AI inference, ஒருங்கிணைக்கப்பட்ட பல-GPU இயக்க pipeline-கள், SRAM-இல் இணைக்கப்பட்ட (fused) Triton/CUDA kernel-கள், மற்றும் மிகக் குறைந்த-latency fabric interconnect-கள் ஆகியவற்றை இணைத்து, மில்லிசெக்கண்டுகளில் token-களை உருவாக்குகிறது.

1. End-to-End Inference Walkthrough: Prompt-இலிருந்து Token வரை

ஒரு பயனாளர் 1,000-token prompt-ஐ, vLLM போன்ற inference engine-இல் Tensor Parallelism பயன்படுத்தி 4 GPU-களில் இயங்கும் 70B-அளவுரு (parameter) LLM-க்கு அனுப்புகிறார் என்று வைத்துக்கொள்வோம்.

• படி 1: Tokenization & Host-to-Device இடமாற்றம் (CPU → GPU).

• CPU tokenizer, உரையை 1,000 முழு எண் token ID-களின் வரிசையாக மாற்றுகிறது.
• Serving engine, PagedAttention வழியே GPU HBM-இல் தொடர்ச்சியற்ற (non-contiguous) இயற்கை நினைவக தொகுதிகளை ஒதுக்கி, token ID-களை PCIe bus வழியே GPU 0-இன் VRAM-க்குள் நகல் செய்கிறது.

• படி 2: Embedding Lookup & ஒலிபரப்பு (Broadcasting) (NVLink).

• GPU 0, 1,000 embedding vector-களை (எ.கா., வடிவம் [1000, 8192]) பார்த்து, அந்த tensor-ஐ NVLink வழியே ஒலிபரப்புகிறது (broadcast); இதனால் 4 GPU-களும் ஒத்த ஆரம்ப activation-களை வைத்திருக்கின்றன.

• படி 3: Prefill கட்டம் (Compute-Bound Tensor Core இயக்கம்).

• 80 Transformer layer-கள் ஒவ்வொன்றுக்குள்ளும், FlashAttention அனைத்து 1,000 prompt token-களுக்கும் ஒரேசமயம் self-attention-ஐக் கணக்கிடுகிறது. விளையும் Key (K) மற்றும் Value (V) tensor-கள் ஒவ்வொரு GPU-இன் PagedAttention KV cache தொகுதிகளிலும் எழுதப்படுகின்றன.
• Feed-Forward Network (FFN) weight-கள் 4 GPU-களில் நிரல் வாரியாகப் (column-wise) பிரிக்கப்பட்டிருப்பதால் (Tensor Parallelism), ஒவ்வொரு GPU-யும் [1000, 8192] activation-களைத் தன் உள்ளூர் weight பகுதியுடன் (slice) FP8/FP4 Tensor Core-களைப் பயன்படுத்திப் பெருக்குகிறது.
• Layer-இன் முடிவில், 4 GPU-களும் NVLink All-Reduce செய்து தமது பகுதி முடிவுகளைக் கூட்டுகின்றன; அடுத்த layer-க்குச் செல்வதற்கு முன் ஒவ்வொரு GPU-யும் முழு புதுப்பிக்கப்பட்ட [1000, 8192] tensor-ஐப் பெறுகிறது.

• படி 4: Sampling & முதல் Token வெளியீடு.

• Layer 80-க்குப் பின், token 1,000-இன் இறுதி hidden state, vocabulary matrix-ஆல் பெருக்கப்பட்டு ~128,000 சொற்களின் மீது logits-ஐ உருவாக்குகிறது. ஒரு sampling kernel முதல் output token-ஐத் தேர்ந்தெடுத்து அதை CPU stream-க்குத் திருப்பி அனுப்புகிறது.

• படி 5: Autoregressive Decode சுழற்சி (CUDA Graphs + Memory-Bound இயக்கம்).

• Token 1,001-ஐ உருவாக்க, CPU நூற்றுக்கணக்கான தனித்தனி kernel-களை மீண்டும் launch செய்யாது. அதற்குப் பதிலாக, ஒரே API அழைப்பில் முன்பதிவு செய்யப்பட்ட CUDA Graph-ஐத் தொடங்கிவிடுகிறது.
• இப்போது ஒரே ஒரு [1, 8192] token vector-ஐ மட்டுமே செயாக்கும் ஒவ்வொரு GPU-யும், தன் உள்ளூர் model weight-களையும் 1,000 cache செய்யப்பட்ட KV vector-களையும் HBM-இலிருந்து on-chip SRAM-க்குள் stream செய்து, புதிய attention வெளியீட்டைக் கணக்கிட்டு, NVLink All-Reduce-ஐ இயக்கி, token 1,001-இன் K மற்றும் V vector-களை அடுத்த இலவச PagedAttention இடத்தில் (slot) சேர்க்கிறது. End-of-Sequence (EOS) token உருவாகும் வரை இச்சுழற்சி தொடர்கிறது.

2. OpenAI Triton-இல் Custom Fused Softmax Kernel கட்டுதல்

ஒரு எளிய (naive) PyTorch Softmax (torch.exp(x) / torch.sum(torch.exp(x))) மூன்று தனித்தனி HBM memory pass-களைத் தேவைப்படுத்துகிறது (எண்ணியல் நிலைத்தன்மைக்காக நிரையின் பெருமத்தைக் (row max) கண்டறிய ஒன்று, exponential-களைக் கணக்கிட்டுக் கூட்ட ஒன்று, வகுக்க ஒன்று).

கீழே OpenAI Triton-இல் முழுமையான, உற்பத்தி-பாணி (production-style) Fused Row-Wise Softmax kernel உள்ளது. இது முழு நிரையையும் வேகமான on-chip SRAM-க்குள் ஒரே ஒரு முறை ஏற்றி, மூன்று படிகளையும் GPU register-களுக்குள் செய்து, இறுதி வெளியீட்டை ஒரே memory pass-இல் 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 Cluster Networking: NVLink, InfiniBand, மற்றும் RoCEv2

AI பணிச்சுமைகளை ஆயிரக்கணக்கான GPU-களில் விரிவாக்கும் (scaling) போது, network latency-யும் நெரிசலும் (congestion) கச்சிதமான GPU FLOPS-ஐப் போலவே முக்கியமாகின்றன. நவீன cluster-கள் இரு-நிலை scale-up மற்றும் scale-out network படிநிலையைப் (hierarchy) பயன்படுத்துகின்றன.

• Collective செயல்பாடுகள் Network-ஐ எவ்வாறு பயன்படுத்துகின்றன.

• All-Reduce (Gradient & Tensor Sync): ஒவ்வொரு GPU-யும் ஒரு பகுதி tensor-உடன் ஆரம்பித்து, cluster முழுவதும் உள்ள அனைத்து tensor-களையும் கூட்டி, ஒத்த இறுதிக் கூட்டுத்தொகையுடன் முடிகிறது. Ring அல்லது Tree வழிமுறைகளைப் — அல்லது switch-கள் switch சிலிக்கானுக்குள்ளேயே vector-களைக் கூட்டும் In-Network Computing (SHARP)-ஐ — பயன்படுத்தி, bandwidth cluster அளவைச் சாராமல் அளவிடப்படுகிறது (scales).
• All-to-All (Mixture-of-Experts Routing): ஒவ்வொரு GPU-யும் அதே தரவை அண்டை GPU-களுக்கு அனுப்பும் All-Reduce-ஐப் போலன்றி, All-to-All ஒவ்வொரு GPU-யும் இலக்கு expert-களை வைத்திருக்கும் மற்ற ஒவ்வொரு GPU-க்கும் தனித்துவமான token துண்டுகளை (slices) ஒரேசமயம் அனுப்பும் அனைத்து-எதிர்-அனைத்து போக்குவரத்து வெடிப்பை உருவாக்குகிறது. இது பெறுநர் port-களில் கடுமையான incast நெரிசலை ஏற்படுத்துகிறது.

• AI Network-களுக்கு இழப்பற்ற (Lossless) RDMA ஏன் தேவை.

• பாரம்பரிய TCP/IP networking, packet-களைச் செயாக்க CPU-ஐ நம்பியிருக்கிறது; buffer-கள் நிரம்பும்போது packet-களைக் கைவிடுகிறது (drops) — இது மில்லிசெக்கண்ட்-அளவு மீளனுப்பல்களைத் (retransmissions) தூண்டி, முழு 10,000-GPU பயிற்சி cluster-ஐயும் நிறுத்திவிடும் (stall).
• InfiniBand, RoCEv2 இரண்டும் RDMA (Remote Direct Memory Access)-ஐப் பயன்படுத்தி, ஒரு GPU-இன் Network Interface Card (NIC) பூச்சிய CPU தலையீட்டுடன் மற்றொரு GPU-இன் HBM VRAM-இல் நேரடியாக வாசிக்கவும்/எழுதவும் அனுமதிக்கிறது. RoCEv2, நிலையான Ethernet-இல் பூச்சிய packet இழப்பை உறுதிசெய்கிறது: switch buffer-கள் நிரம்புவதற்கு முன் மேல்நோக்கு அனுப்புநர்களை இடைநிறுத்த PFC (Priority Flow Control)-ஐயும், அனுப்பும் வீதத்தை மென்மையாகக் குறைக்க (throttle) ECN (Explicit Congestion Notification)-ஐயும் பயன்படுத்துகிறது.