SwiGLU ඒකාබද්ධ කිරීම, FlexAttention සහ GPU ප්‍රශස්තකරණ කෙටි මාර්ගෝපදේශය

ප්‍රායෝගික PyTorch කර්නල සකස් කිරීම, PyTorch FlexAttention සහ පුළුල් ප්‍රශස්තකරණ කෙටි මාර්ගෝපදේශයක් එකට ගෙන ඒමෙන් නවීන GPU ක්‍රමලේඛනය සඳහා ප්‍රායෝගික ක්‍රියාවලිය සම්පූර්ණ වේ. පහත දැක්වෙන්නේ මෙම ක්ෂේත්‍ර තුනම පිළිබඳ ගැඹුරු විමසුමකි.

1. නිශ්චිත PyTorch මෙහෙයුමක් සකස් කිරීම: SwiGLU සක්‍රීයකරණ ඒකාබද්ධ කිරීම

නවීන LLM වල (Llama, Mistral සහ DeepSeek වැනි), Feed-Forward Network (FFN) SwiGLU භාවිත කරයි:

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

eager PyTorch හි, මූලද්‍රව්‍ය-අනුව සක්‍රීයකරණ පියවර F.silu(gate) * up ගණනය කිරීමෙන් කර්නල කිහිපයක් දියත් වන අතර අතරමැදි tensors HBM වෙත ලියයි. ප්‍රශස්තකරණ ස්ථර තුනක් හරහා එය 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

• සකස් කිරීමේ සාරාංශය: @triton.autotune භාවිතයෙන්, Triton ඔබේ නිශ්චිත GPU දෘඩාංග මත විවිධ BLOCK_SIZE සහ num_warps සංයෝජන benchmark කර, HBM කලාප පළල සංතෘප්ත කරන (භෞතික උපරිමයෙන් ~85%–92%) වින්‍යාසය තෝරාගනිමින්, eager PyTorch හා සසඳන විට සක්‍රීයකරණ ප්‍රමාදය 2.5× සිට 3.5× දක්වා අඩු කරයි.

2. PyTorch FlexAttention: Triton හරහා වැඩසටහන්ගත කළ හැකි FlashAttention

ඓතිහාසිකව, පර්යේෂකයන් වේදනාකාරී ද්විත්වයකට මුහුණ දුන්නේය: වේගවත් වේගය සඳහා FlashAttention භාවිත කිරීම (එය සහාය දුන්නේ causal හෝ full attention වැනි දෘඪ mask පමණි) හෝ අභිරුචි මන්දගාමී PyTorch attention mask ලිවීම (එමඟින් VRAM හි ද්විඝාත N × N attention න්‍යාසය ගොඩනැඟුණි).

PyTorch FlexAttention මෙය විසඳන්නේ, සම්මත Python ශ්‍රිත ලිවීමෙන් attention ලකුණු ක්ෂණිකව වෙනස් කිරීමට ඉඩ දීමෙනි; torch.compile එය සෘජුවම ඒකාබද්ධ FlashAttention Triton කර්නලයක් බවට පහත් කරයි.

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 බෑවුම්, සාපේක්ෂ ස්ථානීය නැඹුරුතා, හෝ attention logit soft-capping වැනි) වේගවත් SRAM තුළ FlashAttention අභ්‍යන්තර ලූපයට කෙලින්ම ඇතුළත් කරයි.
• block_mask: attention න්‍යාසයේ 128 × 128 ටයිල් කුමක් සම්පූර්ණයෙන්ම බිංදු කර ඇත්දැයි (උදාහරණ ලෙස, අනුක්‍රම ඇසුරුම් කිරීමේදී ලේඛන සීමාවන් හරහා) කල්තියා ගණනය කර, එම ටයිල් ගණනය කිරීම සම්පූර්ණයෙන්ම මඟ හරින අතර, ඝන causal attention හා සසඳන විට විශාල වේග වැඩිවීම් ලබා දෙයි.

3. ප්‍රධාන කෙටි මාර්ගෝපදේශය: CUDA, Triton සහ LLM ප්‍රශස්තකරණ නීති

බාධක හඳුනාගැනීමේදී හෝ AI පුහුණු සහ අනුමාන (inference) නල මාර්ග නිර්මාණය කිරීමේදී මෙම ඉක්මන් යොමු මාර්ගෝපදේශය භාවිත කරන්න: