SwiGLU fusion, FlexAttention மற்றும் GPU மேம்படுத்தல் (optimization) விரைவுக் குறிப்பேடு

நவீன GPU நிரலாக்கத்திற்கான (GPU programming) நடைமுறைச் செயன்முறையை நிறைவுசெய்யும் வகையில், செயல்முறை அடிப்படையிலான PyTorch kernel சீரமைப்பு (tuning), PyTorch FlexAttention மற்றும் விரிவான மேம்படுத்தல் (optimization) விரைவுக் குறிப்பேடு ஆகியவை இங்கு ஒன்றிணைக்கப்படுகின்றன. இந்த மூன்று துறைகள் பற்றிய விரிவான விளக்கம் கீழே தரப்பட்டுள்ளது.

1. குறிப்பிட்ட PyTorch செயல்பாட்டைச் சீரமைத்தல்: SwiGLU Activation Fusion

நவீன பெரிய மொழி மாதிரிகளில் (LLMs), எடுத்துக்காட்டாக Llama, Mistral மற்றும் DeepSeek ஆகியவற்றில், Feed-Forward Network (FFN), SwiGLU-ஐப் பயன்படுத்துகின்றது:

SwiGLU(x)=SiLU(xW_(gate))\odot (xW_(up))

eager PyTorch-இல், ஒவ்வொரு உறுப்பிற்குமான செயற்படுத்தல் படியான (element-wise activation step) F.silu(gate) * up என்பதைக் கணக்கிடும்போது பல kernels தொடங்கப்படுவதுடன், இடைநிலை tensors, HBM-இல் எழுதப்படுகின்றன. இதனை மூன்று மேம்படுத்தல் (optimization) நிலைகளில் செயல்திறன் அளவீடு (benchmark) செய்து சீரமைக்கும் முறை பின்வருமாறு:

import torch
import torch.nn.functional as F
import triton
import triton.language as tl

# Tier 1: Eager PyTorch Baseline (Multiple HBM read/write passes)
def eager_swiglu(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor:
    return F.silu(gate) * up

# Tier 2: Automated Compiler Fusion via torch.compile
compiled_swiglu = torch.compile(eager_swiglu, mode="max-autotune")

# Tier 3: Custom Autotuned Triton Kernel (Single HBM pass + Block Tuning)
@triton.autotune(
    configs=[
        triton.Config({'BLOCK_SIZE': 1024}, num_warps=4),
        triton.Config({'BLOCK_SIZE': 2048}, num_warps=8),
        triton.Config({'BLOCK_SIZE': 4096}, num_warps=8),
    ],
    key=['n_elements'],
)
@triton.jit
def swiglu_triton_kernel(gate_ptr, up_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements

    # Load both inputs once into registers in FP32 for numerical accuracy
    gate = tl.load(gate_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
    up = tl.load(up_ptr + offsets, mask=mask, other=0.0).to(tl.float32)

    # SiLU(gate) = gate * sigmoid(gate), computed purely in registers
    silu_gate = gate * tl.sigmoid(gate)
    out = (silu_gate * up).to(tl.float16)

    tl.store(out_ptr + offsets, out, mask=mask)

def triton_swiglu(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor:
    out = torch.empty_like(gate)
    n_elements = gate.numel()
    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)
    swiglu_triton_kernel[grid](gate, up, out, n_elements)
    return out

• சீரமைப்பிலிருந்து பெறப்படும் முக்கிய குறிப்பு (Tuning Takeaway): @triton.autotune-ஐப் பயன்படுத்துவதன் மூலம், உங்களுடைய அதே GPU வன்பொருளில் வெவ்வேறு BLOCK_SIZE மற்றும் num_warps சேர்க்கைகளின் செயல்திறனை Triton அளவிடுகின்றது. HBM bandwidth-ஐ முழுமையாகப் பயன்படுத்தும் அமைப்பைத் தெரிவுசெய்து, அதாவது வன்பொருளின் உச்ச வரம்பில் சுமார் ~85%–92%-ஐ எட்டும் அமைப்பைத் தெரிவுசெய்து, eager PyTorch உடன் ஒப்பிடுகையில் activation latency-ஐ 2.5× முதல் 3.5× வரை குறைக்கின்றது.

2. PyTorch FlexAttention: Triton ஊடாக நிரலாக்கக்கூடிய FlashAttention

முன்னர் ஆய்வாளர்கள் கடினமான இரு தெரிவுகளுக்கு இடையில் சிக்கியிருந்தனர்: வேகத்திற்காக FlashAttention-ஐப் பயன்படுத்துவது, ஆனால் அது causal அல்லது full attention போன்ற மாற்றத்திற்குக் கட்டுப்பாடுள்ள masks-ஐ மட்டுமே ஆதரித்தது; அல்லது தேவைக்கேற்ப மெதுவான PyTorch attention masks-ஐ எழுதுவது, ஆனால் அவை இருபடிச் சார்புடைய (quadratic) N × N attention matrix-ஐ VRAM-இல் முழுமையாக உருவாக்கிச் சேமித்தன.

செயல்பாட்டின்போதே attention scores-ஐ மாற்றுவதற்கு வழமையான Python functions-ஐ எழுத அனுமதிப்பதன் மூலம் PyTorch FlexAttention இந்தச் சிக்கலைத் தீர்க்கின்றது. அவற்றை torch.compile நேரடியாக ஒன்றிணைக்கப்பட்ட FlashAttention Triton kernel ஆக மாற்றுகின்றது.

import torch
from torch.nn.attention.flex_attention import flex_attention, create_block_mask

# Example 1: Custom Score Modification (ALiBi Positional Bias + Soft-Capping)
def alibi_softcap_score_mod(score, b, h, q_idx, kv_idx):
    # Apply Gemma-style logit soft-capping and distance penalty inside SRAM
    score = 30.0 * torch.tanh(score / 30.0)
    distance = q_idx - kv_idx
    return score - (0.125 * (h + 1) * distance)

# Example 2: Sparse Block Mask (Sliding Window + Document Packing Isolation)
document_id = torch.tensor([0, 0, 0, 1, 1, 1, 1, 1], device="cuda")

def doc_sliding_window_mask(b, h, q_idx, kv_idx):
    causal = q_idx >= kv_idx
    window = (q_idx - kv_idx) <= 512
    same_doc = document_id[q_idx] == document_id[kv_idx]
    return causal & window & same_doc

# Pre-compile sparse block mask so completely masked 128x128 tiles are skipped entirely
block_mask = create_block_mask(doc_sliding_window_mask, B=1, H=8, Q_LEN=8, KV_LEN=8)
compiled_flex_attention = torch.compile(flex_attention)

# Executes at FlashAttention speeds with zero N x N VRAM materialization
# out = compiled_flex_attention(q, k, v, score_mod=alibi_softcap_score_mod, block_mask=block_mask)

• FlexAttention ஒரு முக்கிய முன்னேற்றமாக இருப்பது ஏன்:

• score_mod: ALiBi slopes, relative positional biases அல்லது attention logit soft-capping போன்ற, தேவைக்கேற்ப அமைக்கப்பட்ட ஒவ்வொரு உறுப்பிற்குமான கணிதச் செயல்பாடுகளை (element-wise math), வேகமான SRAM-இல் இயங்கும் FlashAttention inner loop-இற்குள் நேரடியாகச் சேர்க்கின்றது.
• block_mask: attention matrix-இன் எந்த 128 × 128 tiles முற்றாகப் பூச்சியமாக்கப்பட்டுள்ளன என்பதை முன்கூட்டியே கணக்கிடுகின்றது. எடுத்துக்காட்டாக, sequence packing-இன்போது ஆவணங்களுக்கு இடையிலான எல்லைகள் இவ்வாறு அமையலாம். அந்த tiles-இற்கான கணக்கீட்டை முற்றாகத் தவிர்ப்பதால், dense causal attention உடன் ஒப்பிடுகையில் மிகப்பெரிய வேக அதிகரிப்பைப் பெற முடிகின்றது.

3. முதன்மை விரைவுக் குறிப்பேடு (Master Cheat-Sheet): CUDA, Triton மற்றும் LLM மேம்படுத்தல் (Optimization) விதிகள்

செயல்திறன் தடைகளை (bottlenecks) கண்டறியும்போதோ, AI பயிற்சி (training) மற்றும் inference pipelines-இன் கட்டமைப்பை வடிவமைக்கும்போதோ இந்த விரைவுக் குறிப்பு வழிகாட்டியைப் பயன்படுத்தவும்: