افزایش سرعت توجه در پردازنده گرافیکی (GPU) با cuDNN و TransformerEngine

۱. مقدمه

جکس در مسیر یادگیری GPU. آزمایشگاه ۵: توجه به GPU.

در این آزمایشگاه کد، شما بر محاسبه توجه در پردازنده‌های گرافیکی NVIDIA تمرکز می‌کنید. شما با یک پیاده‌سازی ساده شروع می‌کنید، آن را با توجه ترکیبی داخلی JAX جایگزین می‌کنید، سپس cuDNN NVIDIA را برای همان عملیات فعال می‌کنید. در پایان، شما خواهید دانست که در هر سطح چه تغییراتی رخ می‌دهد و مسیرهای بهینه شده برای GPU با رشد مسئله چقدر سریع‌تر می‌شوند.

کاری که انجام خواهید داد

  • پیاده‌سازی توجه نقطه‌ای مقیاس‌پذیر از ابتدا با عملیات پایه JAX
  • آن را با jax.nn.dot_product_attention ، هسته فیوز شده داخلی JAX، جایگزین کنید.
  • بک‌اند cuDNN را با implementation="cudnn" اجباری کنید و پوشش سببی را اضافه کنید
  • طول توالی و اندازه دسته را جابجا کنید تا ببینید هسته‌های ادغام‌شده در کجا نتیجه می‌دهند.
  • شکل‌های توجه چند سر MHA، GQA و MQA را با هم مقایسه کنید
  • توجه به NVIDIA TransformerEngine را محک بزنید و مسیر کد FP8 آن را بررسی کنید

آنچه نیاز دارید

  • یک پروژه Google Cloud با قابلیت پرداخت صورتحساب، و اعتبار کارگاه یا رزرو شامل استفاده از GPU
  • سهمیه حداقل ۲ پردازنده گرافیکی NVIDIA L4 در منطقه انتخابی شما ( نحوه بررسی سهمیه پردازنده گرافیکی )
  • تکمیل آزمایشگاه‌های کد ۱ تا ۴، یا یک محیط معادل JAX GPU مجهز به CUDA
  • آشنایی با ضرب ماتریس و softmax. نیازی به تجربه CUDA یا cuDNN نیست.

زمان تخمینی برای تکمیل: ۶۰ دقیقه .

نحوه عملکرد توجه

توجه ضرب نقطه‌ای مقیاس‌پذیر سه ورودی - پرس‌وجو ، کلید و مقدار - را دریافت کرده و موارد زیر را محاسبه می‌کند:

Attention(Q, K, V) = softmax(Q K^T / sqrt(d_k)) V

ضریب مقیاس‌بندی 1 / sqrt(d_k) مانع از بزرگ شدن بیش از حد حاصلضرب‌های نقطه‌ای با افزایش ابعاد سر می‌شود، که این امر softmax را به مناطقی سوق می‌دهد که گرادیان‌های آن کوچک هستند.

چهار مرحله، که هر کدام مقدمه‌ی مرحله‌ی بعدی است:

  1. امتیاز — Q @ KT : حاصلضرب نقطه‌ای نشان می‌دهد که هر موقعیت جستجو چقدر باید به هر موقعیت کلیدی توجه کند.
  2. مقیاس + Softmax — softmax(scores / sqrt(d)) : مقیاس‌بندی از محو شدن گرادیان جلوگیری می‌کند و softmax نمرات را به وزن‌های توجه تبدیل می‌کند که مجموع آنها برابر با ۱ می‌شود.
  3. حضور — weights @ V : مجموع وزن‌دار بردارهای ارزش، خروجی را برای هر موقعیت پرس‌وجو تولید می‌کند.
  4. خروجی — همان شکل Q: هر موقعیت پرس‌وجو اکنون اطلاعاتی از موقعیت‌هایی که به آنها توجه کرده است را در خود جای می‌دهد.

۲. قبل از شروع

پروژه خود را انتخاب کنید

در کنسول گوگل کلود ، یک پروژه با قابلیت پرداخت فعال انتخاب یا ایجاد کنید.

پوسته ابری را باز کنید

برای شروع یک جلسه Cloud Shell ، روی Activate Cloud Shell (آیکون ترمینال در سمت راست بالای کنسول) کلیک کنید، سپس آن را به پروژه خود هدایت کنید:

gcloud config set project <YOUR_PROJECT_ID>

فراهم کردن محیط GPU

دستور زیر را در Cloud Shell اجرا کنید. Codelab 1 هر دستور را با جزئیات توضیح می‌دهد، از جمله سهمیه GPU و الزامات منطقه.

gcloud services enable \
  container.googleapis.com \
  compute.googleapis.com \
  iam.googleapis.com \
  cloudresourcemanager.googleapis.com \
  logging.googleapis.com \
  monitoring.googleapis.com

git clone https://github.com/Google-Cloud-AI/partner-ai-nvidia.git
cd partner-ai-nvidia/05-workshops/jax-on-gpu/terraform

cp terraform.tfvars.example terraform.tfvars

فایل terraform.tfvars را ویرایش کنید و project_id = " " قرار دهید. project_id = " " ، سپس کلاستر را آماده کرده و JupyterLab را مستقر کنید:

terraform init
terraform apply
$(terraform output -raw get_credentials_command)

cd ..
kubectl apply -f deploy/jupyter.yaml

terraform apply حدود ۱۲ دقیقه طول می‌کشد. پس از اتمام، منتظر Pod و LoadBalancer باشید، سپس توکن یکبار مصرف JupyterLab را از لاگ Pod بخوانید:

kubectl get pod jax-jupyter -w         # wait for Running, then Ctrl+C
kubectl get svc jax-jupyter-svc -w     # wait for EXTERNAL-IP, then Ctrl+C
kubectl logs jax-jupyter | grep -o 'token=[a-z0-9]*' | head -1

http:// :8884 را باز کنید http:// :8884 وارد کنید، توکن را جایگذاری کنید و یک دفترچه یادداشت پایتون ۳ جدید در /workspace ایجاد کنید. هر بلوک کد در این codelab به سلولی از آن دفترچه یادداشت می‌رود.

آنچه این codelab نیاز دارد را نصب کنید 

!pip install --quiet matplotlib flax

GPU را تنظیم و تأیید کنید

JAX را وارد کنید و تأیید کنید که backend پیش‌فرض یک GPU است. این سلول همچنین block_tree ، show_table و show_bars را تعریف می‌کند، کمکی‌هایی که در هر مرحله بعدی برای انتظار برای کار دستگاه و ارائه نتایج استفاده می‌شوند.

import os
os.environ["LD_LIBRARY_PATH"] = "/usr/local/nvidia/lib64:" + os.environ.get("LD_LIBRARY_PATH", "")

import html
import math
import time
from functools import partial

from IPython.display import HTML, display
import matplotlib.pyplot as plt
import numpy as np

import jax
import jax.numpy as jnp


devices = jax.devices()
gpu_devices = [d for d in devices if d.platform == "gpu"]
device = gpu_devices[0] if gpu_devices else None

print(f"JAX version:     {jax.__version__}")
print(f"Default backend: {jax.default_backend()}")
print(f"Devices:         {devices}")

assert gpu_devices, f"This lab assumes a GPU backend. Available devices: {devices}"
print(f"Using GPU:       {device}")


def block_tree(tree):
    """Wait until a PyTree of JAX arrays is ready on device."""
    return jax.block_until_ready(tree)


def show_table(headers, rows, title=None, aligns=None):
    """Render rows as an HTML table."""
    aligns = aligns or ["left"] * len(headers)
    parts = ["<div style='font-family: system-ui; max-width: 980px;'>"]
    if title:
        parts.append(f"<h4 style='margin: 0 0 8px 0;'>{html.escape(title)}</h4>")
    parts.append("<table style='border-collapse: collapse; width: 100%; font-size: 13px;'>")
    parts.append("<thead><tr>")
    for h, a in zip(headers, aligns):
        parts.append(
            f"<th style='text-align:{a}; border-bottom:1px solid #d0d7de; padding:6px;'>"
            f"{html.escape(str(h))}</th>"
        )
    parts.append("</tr></thead><tbody>")
    for row in rows:
        parts.append("<tr>")
        for cell, a in zip(row, aligns):
            parts.append(
                f"<td style='text-align:{a}; border-bottom:1px solid #eef1f4; padding:6px;'>"
                f"{html.escape(str(cell))}</td>"
            )
        parts.append("</tr>")
    parts.append("</tbody></table></div>")
    display(HTML("".join(parts)))


def show_bars(rows, title, unit="", lower_is_better=False):
    """Render (label, value) pairs as a horizontal bar chart in HTML."""
    max_value = max(float(value) for _, value in rows) or 1.0
    color = "#1a7f37" if not lower_is_better else "#0969da"
    parts = ["<div style='font-family: Arial, sans-serif; max-width: 760px;'>"]
    parts.append(f"<h4 style='margin: 0 0 8px 0;'>{html.escape(title)}</h4>")
    for label, value in rows:
        width = max(3, 100 * float(value) / max_value)
        parts.append(
            "<div style='display:grid; grid-template-columns: 190px 1fr 130px; gap: 8px; "
            "align-items:center; margin: 6px 0;'>"
            f"<div style='font-size:13px;'>{html.escape(str(label))}</div>"
            "<div style='background:#f6f8fa; border-radius:6px; overflow:hidden; height:22px;'>"
            f"<div style='height:22px; width:{width:.1f}%; background:{color};'></div></div>"
            f"<div style='font-size:13px; font-variant-numeric: tabular-nums;'>{float(value):,.1f} {html.escape(unit)}</div>"
            "</div>"
        )
    parts.append(
        f"<div style='font-size:12px; color:#57606a;'>"
        f"{'Lower' if lower_is_better else 'Higher'} is better.</div></div>"
    )
    display(HTML("".join(parts)))

شما باید نسخه JAX، gpu به عنوان backend پیش‌فرض، لیستی از دستگاه‌های CUDA و پردازنده گرافیکی (GPU) مورد استفاده در بقیه codelab را ببینید.

۳. ساخت آرایه‌های آزمایشی Q، K و V

هسته‌های توجه ادغام‌شده به نحوه چیدمان آرایه‌های ورودی در حافظه اهمیت می‌دهند، بنابراین قبل از نوشتن هرگونه کد توجه، Q، K و V را در چیدمانی که JAX انتظار دارد، بسازید. jax.nn.dot_product_attention انتظار چیدمان زیر را دارد:

کم نور

معنی

پیش‌فرض ما

ب

اندازه دسته

۴

تی

طول توالی پرس و جو

۱۲۸

س

طول توالی کلید/مقدار

۱۲۸ (همانند T برای توجه به خود)

ن

تعداد سرهای توجه

۸

ح

ابعاد سرانه

۶۴

برای پیاده‌سازی ساده، از float32 شروع می‌کنید و با مقادیر تصادفی مقداردهی اولیه می‌کنید.

BATCH = 4
SEQ_LEN = 128
NUM_HEADS = 8
HEAD_DIM = 64

key = jax.random.key(0)
k1, k2, k3 = jax.random.split(key, 3)

q = jax.random.normal(k1, (BATCH, SEQ_LEN, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
k = jax.random.normal(k2, (BATCH, SEQ_LEN, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
v = jax.random.normal(k3, (BATCH, SEQ_LEN, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)

q, k, v = jax.device_put((q, k, v), device)

show_table(
    ["Array", "Shape", "Dtype", "Layout"],
    [
        ("Q (query)", q.shape, q.dtype, "(B, T, N, H)"),
        ("K (key)", k.shape, k.dtype, "(B, S, N, H)"),
        ("V (value)", v.shape, v.dtype, "(B, S, N, H)"),
    ],
    title="Attention inputs on GPU",
)

شما باید جدولی با یک ردیف برای هر آرایه ببینید، که هر کدام شکل (4, 128, 8, 64) و نوع داده float32 را گزارش می‌دهند. فراخوانی jax.device_put هر سه آرایه را به GPU انتخاب شده در سلول تنظیمات پین می‌کند، بنابراین هیچ چیز در معیارهای بعدی، انتقال میزبان را اندازه‌گیری نمی‌کند.

۴. توجه را از ابتدا پیاده‌سازی کنید

قبل از اینکه به سراغ یک هسته ترکیبی بروید، فرمول را مستقیماً بنویسید تا بتوانید دقیقاً ببینید هسته جایگزین چه چیزی می‌شود. این پیاده‌سازی به ترتیب از چهار مرحله پیروی می‌کند: محاسبه امتیازها، مقیاس‌بندی، softmax، ضرب در مقادیر. شما ابعاد سر و توالی را ترانهاده می‌کنید تا matmul روی محور توالی عمل کند.

این روش به درستی کار می‌کند، اما سه راه‌اندازی جداگانه هسته پردازنده گرافیکی برای ماتریس امتیاز، سافت‌مکس و ماتریس مقدار انجام می‌دهد و ماتریس وزن توجه کامل (B, N, T, S) را در حافظه پردازنده گرافیکی محقق می‌کند.

def naive_attention(q, k, v):
    """Scaled dot-product attention from scratch."""
    scale = 1.0 / math.sqrt(q.shape[-1])

    # (B, T, N, H) to (B, N, T, H) so matmul runs over the T/S axis
    q_t = jnp.transpose(q, (0, 2, 1, 3))
    k_t = jnp.transpose(k, (0, 2, 1, 3))
    v_t = jnp.transpose(v, (0, 2, 1, 3))

    # Score: (B, N, T, H) @ (B, N, H, S) to (B, N, T, S)
    scores = jnp.matmul(q_t, jnp.transpose(k_t, (0, 1, 3, 2))) * scale
    weights = jax.nn.softmax(scores, axis=-1)

    # Attend: (B, N, T, S) @ (B, N, S, H) to (B, N, T, H)
    out_t = jnp.matmul(weights, v_t)

    # Back to (B, T, N, H)
    return jnp.transpose(out_t, (0, 2, 1, 3))


naive_out = block_tree(naive_attention(q, k, v))

show_table(
    ["", "Value"],
    [
        ("Output shape", str(naive_out.shape)),
        ("Output dtype", str(naive_out.dtype)),
    ],
    title="Naive attention",
)

شما باید شکل خروجی (4, 128, 8, 64) و dtype float32 را ببینید که همان شکل Q است، که همان چیزی است که مرحله خروجی فرمول وعده داده بود.

۵. به dot_product_attention بروید

تابع jax.nn.dot_product_attention مراحل امتیاز، مقیاس، softmax و توجه را در یک عملیات واحد ادغام می‌کند. سپس JAX و XLA می‌توانند الگوی دسترسی به حافظه را بهینه کنند. به طور خاص، آنها می‌توانند از تحقق کامل ماتریس وزن توجه در زمانی که دنباله طولانی است، جلوگیری کنند.

با implementation=None ، JAX بهترین backend موجود را به طور خودکار انتخاب می‌کند. در یک GPU با cuDNN موجود و ورودی‌های سازگار، ممکن است از cuDNN استفاده کند. در سخت‌افزارهای دیگر، به XLA برمی‌گردد. 

sdpa_out = block_tree(jax.nn.dot_product_attention(q, k, v))

max_diff = float(jnp.max(jnp.abs(naive_out - sdpa_out)))

show_table(
    ["", "Value"],
    [
        ("Output shape", str(sdpa_out.shape)),
        ("Output dtype", str(sdpa_out.dtype)),
        ("Max |naive − SDPA|", f"{max_diff:.2e}"),
        ("Outputs close (atol=1e-3)", str(bool(jnp.allclose(naive_out, sdpa_out, atol=1e-3)))),
    ],
    title="JAX SDPA vs naive",
)

شما باید همان شکل و نوع داده (dtype) نسخه ساده (naive) را ببینید، یک اختلاف مطلق حداکثر کوچک، و True برای بررسی نزدیکی atol=1e-3 .

۶. بک‌اند ترکیبی cuDNN را اجباری کنید

اجازه دادن به JAX برای انتخاب backend راحت است، اما از بیرون نمی‌توانید تشخیص دهید که کدام هسته واقعاً اجرا شده است. تنظیم implementation="cudnn" هسته‌های توجه ترکیبی cuDNN انویدیا را مجبور می‌کند. اینها هسته‌های GPU بهینه شده دستی هستند که کل محاسبه توجه را در یک اجرای هسته واحد با الگوهای دسترسی به حافظه بهینه شده ترکیب می‌کنند.

توجه داشته باشید که cuDNN fused attention الزامات سخت‌افزاری و شکلی مانند قابلیت محاسبه GPU دارد که باید >= 8.0 (Ampere یا جدیدتر) یا float16 یا bfloat16 به عنوان dtypes ورودی باشد. اگر الزامات برآورده نشوند و شما implementation="cudnn" تنظیم کنید، JAX به جای اینکه به عقب برگردد، خطایی ایجاد می‌کند. به همین دلیل است که کد زیر به bfloat16 تبدیل می‌شود و فراخوانی را در try / except قرار می‌دهد: HAS_CUDNN_SDPA را تنظیم می‌کند تا در هر مرحله بعدی بداند که آیا مسیر cuDNN در این دستگاه موجود است یا خیر. 

q_bf16 = q.astype(jnp.bfloat16)
k_bf16 = k.astype(jnp.bfloat16)
v_bf16 = v.astype(jnp.bfloat16)

HAS_CUDNN_SDPA = False

try:
    cudnn_out = block_tree(
        jax.nn.dot_product_attention(q_bf16, k_bf16, v_bf16, implementation="cudnn")
    )
    HAS_CUDNN_SDPA = True

    xla_bf16_out = block_tree(
        jax.nn.dot_product_attention(q_bf16, k_bf16, v_bf16, implementation="xla")
    )
    max_diff = float(jnp.max(jnp.abs(
        cudnn_out.astype(jnp.float32) - xla_bf16_out.astype(jnp.float32)
    )))

    show_table(
        ["", "Value"],
        [
            ("Output shape", str(cudnn_out.shape)),
            ("Output dtype", str(cudnn_out.dtype)),
            ("Max |cuDNN − XLA| (both bf16)", f"{max_diff:.2e}"),
            ("Outputs close (rtol=1e-2, atol=1e-2)", str(bool(jnp.allclose(cudnn_out, xla_bf16_out, rtol=1e-2, atol=1e-2)))),
        ],
        title="cuDNN fused attention",
    )

except Exception as e:
    print(f"cuDNN SDPA not available on this GPU: {e}")
    print("Continuing with XLA backend only.")

اگر پردازنده گرافیکی شما نتواند cuDNN SDPA را اجرا کند، کد دلیل آن را چاپ می‌کند و codelab در بک‌اند XLA ادامه می‌یابد.

۷. پوشش علی را اضافه کنید

در مدل‌های خودهمبسته (دیکودرهای سبک GPT)، هر موقعیت فقط می‌تواند به موقعیت‌های قبلی توجه کند. تنظیم is_causal=True این ماسک مثلثی-پایین را درون هسته ادغام‌شده اعمال می‌کند، بنابراین نیازی نیست خودتان یک ماتریس ماسک بسازید.

کد زیر دو بار، با ماسک و بدون ماسک، توجه را اجرا می‌کند و دو موقعیت را مقایسه می‌کند تا نشان دهد ماسک چه چیزی را تغییر داده است. 

causal_out = block_tree(
    jax.nn.dot_product_attention(q, k, v, is_causal=True)
)

# With causal masking, the last position attends to all positions.
# The first position attends only to itself.
nocausal_out = block_tree(
    jax.nn.dot_product_attention(q, k, v, is_causal=False)
)

# First position should differ
first_pos_diff = float(jnp.max(jnp.abs(causal_out[:, 0] - nocausal_out[:, 0])))
# Last position should be the same
last_pos_diff = float(jnp.max(jnp.abs(causal_out[:, -1] - nocausal_out[:, -1])))

show_table(
    ["Position", "Max diff (causal vs full)", "Expected"],
    [
        ("First (t=0)", f"{first_pos_diff:.4f}", "Large — causal restricts to self only"),
        ("Last (t=T-1)", f"{last_pos_diff:.2e}", "~0 — attends to all positions either way"),
    ],
    title="Causal masking effect on attention output",
)

شما باید در موقعیت اول تفاوت زیادی و در موقعیت آخر تفاوت تقریباً صفر را ببینید. این عدم تقارن به این صورت است که ماسک با موقعیت ۰ هر کلید به جز کلید خودش را از دست می‌دهد، در حالی که موقعیت نهایی می‌توانست کل دنباله را ببیند.

۸. اجرای برخی بنچمارک‌ها

اکنون شما سه روش برای محاسبه‌ی همان تابع دارید. اینکه به کدام یک از آنها برسید کاملاً به شکل مسئله بستگی دارد و می‌توانید آن را ارزیابی کنید.

زمان‌بندی متغیرها در شکل پیش‌فرض 

def benchmark_attention(fn, q, k, v, warmup=3, repeats=50):
    """Time an attention function. Returns median milliseconds per call."""
    jit_fn = jax.jit(fn)

    for _ in range(warmup):
        block_tree(jit_fn(q, k, v))

    times = []
    for _ in range(repeats):
        start = time.perf_counter()
        block_tree(jit_fn(q, k, v))
        times.append((time.perf_counter() - start) * 1000)

    return np.median(times)


t_naive = benchmark_attention(naive_attention, q, k, v)
t_sdpa = benchmark_attention(
    lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="xla"),
    q, k, v,
)

results = [
    ("Naive (matmul + softmax + matmul)", f"{t_naive:.2f}"),
    ("SDPA (XLA, float32)", f"{t_sdpa:.2f}"),
]
bar_data = [
    ("Naive", t_naive),
    ("SDPA XLA f32", t_sdpa),
]

if HAS_CUDNN_SDPA:
    t_sdpa_bf16 = benchmark_attention(
        lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="xla"),
        q_bf16, k_bf16, v_bf16,
    )
    t_cudnn = benchmark_attention(
        lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="cudnn"),
        q_bf16, k_bf16, v_bf16,
    )
    results.append(("SDPA (XLA, bfloat16)", f"{t_sdpa_bf16:.2f}"))
    results.append(("SDPA (cuDNN, bfloat16)", f"{t_cudnn:.2f}"))
    bar_data.append(("SDPA XLA bf16", t_sdpa_bf16))
    bar_data.append(("SDPA cuDNN bf16", t_cudnn))

show_table(
    ["Implementation", "Median ms/call"],
    results,
    title=f"Attention timing — B={BATCH}, T={SEQ_LEN}, N={NUM_HEADS}, H={HEAD_DIM}",
    aligns=["left", "right"],
)
show_bars(bar_data, "Attention latency (ms per call)", "ms", lower_is_better=True)

برای این اندازه کوچک مسئله، تمام پیاده‌سازی‌ها بسیار سریع و عمدتاً محدود به سربار هستند، بنابراین نسخه ساده رقابتی است. cuDNN bf16 باید کمی سریع‌تر باشد، اما مزیت توجه ترکیبی معمولاً در طول توالی‌های طولانی‌تر، که در آن اجتناب از ماتریس توجه کامل اهمیت دارد، آشکارتر می‌شود.

طول دنباله را جارو کنید

مزیت هسته‌های توجه ادغام‌شده با طول توالی افزایش می‌یابد. پیاده‌سازی ساده، یک ماتریس توجه (B, N, T, S) را در حافظه GPU که حافظه O(T²) است، محقق می‌کند. هسته‌های ادغام‌شده مانند cuDNN FlashAttention محاسبات را کاشی‌کاری می‌کنند، بنابراین هرگز ماتریس کامل را محقق نمی‌کنند و حافظه را در O(T) نگه می‌دارند.

رفت و برگشت زیر، هر پیاده‌سازی را در طول توالی از ۶۴ تا ۱۰۲۴ ضرب می‌کند. 

SEQ_LENS = [64, 128, 256, 512, 1024]

sweep_results = []
for sl in SEQ_LENS:
    rk = jax.random.key(sl)
    rk1, rk2, rk3 = jax.random.split(rk, 3)

    q_s = jax.random.normal(rk1, (BATCH, sl, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
    k_s = jax.random.normal(rk2, (BATCH, sl, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
    v_s = jax.random.normal(rk3, (BATCH, sl, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
    q_s, k_s, v_s = jax.device_put((q_s, k_s, v_s), device)

    q_sb = q_s.astype(jnp.bfloat16)
    k_sb = k_s.astype(jnp.bfloat16)
    v_sb = v_s.astype(jnp.bfloat16)

    row = {"seq_len": sl}

    row["naive_ms"] = benchmark_attention(
        naive_attention,
        q_s, k_s, v_s,
        warmup=2,
        repeats=20,
    )

    row["sdpa_xla_f32_ms"] = benchmark_attention(
        lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="xla"),
        q_s, k_s, v_s,
        warmup=2,
        repeats=20,
    )

    row["sdpa_xla_bf16_ms"] = benchmark_attention(
        lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="xla"),
        q_sb, k_sb, v_sb,
        warmup=2,
        repeats=20,
    )

    if HAS_CUDNN_SDPA:
        row["cudnn_bf16_ms"] = benchmark_attention(
            lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="cudnn"),
            q_sb, k_sb, v_sb,
            warmup=2,
            repeats=20,
        )

    sweep_results.append(row)


headers = ["Seq len", "Naive (ms)", "SDPA XLA f32 (ms)", "SDPA XLA bf16 (ms)"]
if HAS_CUDNN_SDPA:
    headers.append("cuDNN bf16 (ms)")

table_rows = []
for r in sweep_results:
    row = [
        r["seq_len"],
        f"{r['naive_ms']:.2f}",
        f"{r['sdpa_xla_f32_ms']:.2f}",
        f"{r['sdpa_xla_bf16_ms']:.2f}",
    ]

    if HAS_CUDNN_SDPA:
        row.append(f"{r['cudnn_bf16_ms']:.2f}")

    table_rows.append(row)

show_table(
    headers,
    table_rows,
    title=f"Sequence-length sweep — B={BATCH}, N={NUM_HEADS}, H={HEAD_DIM}",
    aligns=["right"] * len(headers),
)

عملیات جارو کردن کمی طول می‌کشد، زیرا هر طول توالی قبل از زمان‌بندی، سه یا چهار پیاده‌سازی جداگانه را کامپایل می‌کند. در نهایت باید به ازای هر طول توالی از ۶۴ تا ۱۰۲۴، یک ردیف جدول داشته باشید.

این نمودار، چگونگی تغییر تأخیر توجه با طول توالی را برای پیاده‌سازی ساده، JAX/XLA SDPA در float32 و bfloat16، و cuDNN توجه ترکیبی در bfloat16 مقایسه می‌کند. 

fig, ax = plt.subplots(figsize=(8, 5))
seq_lens = [r["seq_len"] for r in sweep_results]

ax.plot(
    seq_lens,
    [r["naive_ms"] for r in sweep_results],
    "o-",
    label="Naive",
    color="#d1242f",
)

ax.plot(
    seq_lens,
    [r["sdpa_xla_f32_ms"] for r in sweep_results],
    "s-",
    label="SDPA XLA f32",
    color="#0969da",
)

ax.plot(
    seq_lens,
    [r["sdpa_xla_bf16_ms"] for r in sweep_results],
    "d-",
    label="SDPA XLA bf16",
    color="#8250df",
)

if HAS_CUDNN_SDPA:
    ax.plot(
        seq_lens,
        [r["cudnn_bf16_ms"] for r in sweep_results],
        "^-",
        label="cuDNN bf16",
        color="#1a7f37",
    )

ax.set_xlabel("Sequence length")
ax.set_ylabel("Median ms per call")
ax.set_title("Attention latency vs sequence length")
ax.legend()
ax.grid(True, alpha=0.25)
ax.set_xticks(seq_lens)

fig.tight_layout()
plt.show()

شما باید به ازای هر پیاده‌سازی یک خط ببینید. و باید فاصله‌ای را که با طولانی‌تر شدن دنباله افزایش می‌یابد، ببینید.

اندازه دسته را جابجا کنید

دسته‌های بزرگتر، سربار اجرای هسته را کاهش داده و استفاده از پردازنده گرافیکی را تا جایی که حافظه پردازنده گرافیکی به گلوگاه تبدیل شود، بهبود می‌بخشند. جابجایی زیر طول توالی را روی ۲۵۶ ثابت نگه می‌دارد و اندازه دسته را تغییر می‌دهد. 

BATCH_SIZES = [1, 2, 4, 8, 16]
SWEEP_SEQ = 256

batch_results = []
for bs in BATCH_SIZES:
    rk = jax.random.key(bs + 100)
    rk1, rk2, rk3 = jax.random.split(rk, 3)

    q_b = jax.random.normal(rk1, (bs, SWEEP_SEQ, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
    k_b = jax.random.normal(rk2, (bs, SWEEP_SEQ, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
    v_b = jax.random.normal(rk3, (bs, SWEEP_SEQ, NUM_HEADS, HEAD_DIM), dtype=jnp.float32)
    q_b, k_b, v_b = jax.device_put((q_b, k_b, v_b), device)

    q_bb = q_b.astype(jnp.bfloat16)
    k_bb = k_b.astype(jnp.bfloat16)
    v_bb = v_b.astype(jnp.bfloat16)

    row = {"batch": bs}

    row["sdpa_xla_f32_ms"] = benchmark_attention(
        lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="xla"),
        q_b, k_b, v_b,
        warmup=2,
        repeats=20,
    )

    row["sdpa_xla_bf16_ms"] = benchmark_attention(
        lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="xla"),
        q_bb, k_bb, v_bb,
        warmup=2,
        repeats=20,
    )

    if HAS_CUDNN_SDPA:
        row["cudnn_bf16_ms"] = benchmark_attention(
            lambda q, k, v: jax.nn.dot_product_attention(q, k, v, implementation="cudnn"),
            q_bb, k_bb, v_bb,
            warmup=2,
            repeats=20,
        )

    batch_results.append(row)


headers = ["Batch size", "SDPA XLA f32 (ms)", "SDPA XLA bf16 (ms)"]
if HAS_CUDNN_SDPA:
    headers.append("cuDNN bf16 (ms)")

table_rows = []
for r in batch_results:
    row = [
        r["batch"],
        f"{r['sdpa_xla_f32_ms']:.2f}",
        f"{r['sdpa_xla_bf16_ms']:.2f}",
    ]

    if HAS_CUDNN_SDPA:
        row.append(f"{r['cudnn_bf16_ms']:.2f}")

    table_rows.append(row)

show_table(
    headers,
    table_rows,
    title=f"Batch-size sweep — T={SWEEP_SEQ}, N={NUM_HEADS}, H={HEAD_DIM}",
    aligns=["right"] * len(headers),
)


fig, ax = plt.subplots(figsize=(8, 5))
batches = [r["batch"] for r in batch_results]

ax.plot(
    batches,
    [r["sdpa_xla_f32_ms"] for r in batch_results],
    "s-",
    label="SDPA XLA f32",
    color="#0969da",
)

ax.plot(
    batches,
    [r["sdpa_xla_bf16_ms"] for r in batch_results],
    "d-",
    label="SDPA XLA bf16",
    color="#8250df",
)

if HAS_CUDNN_SDPA:
    ax.plot(
        batches,
        [r["cudnn_bf16_ms"] for r in batch_results],
        "^-",
        label="cuDNN bf16",
        color="#1a7f37",
    )

ax.set_xlabel("Batch size")
ax.set_ylabel("Median ms per call")
ax.set_title("Attention latency vs batch size")
ax.legend()
ax.grid(True, alpha=0.25)
ax.set_xticks(batches)

fig.tight_layout()
plt.show()

شما باید یک جدول و یک نمودار با یک ورودی به ازای هر اندازه دسته از ۱ تا ۱۶ دریافت کنید. پیاده‌سازی ساده از این پیمایش حذف شده است، بنابراین هر دو خروجی فقط سه مسیر ادغام‌شده را مقایسه می‌کنند.

۹. MHA، GQA و MQA را با هم مقایسه کنید

تاکنون هر هد کلیدها و مقادیر خاص خود را حمل می‌کرده است. اما رمزگشاها اغلب این تقارن را می‌شکنند، زیرا در طول استنتاج، حافظه پنهان کلید-مقدار چیزی است که بر حافظه تسلط دارد.

توجه چند سر (MHA) به هر سر، پیش‌بینی‌های Q، K و V مخصوص به خود را می‌دهد. توجه پرس‌وجوی گروهی (GQA) و توجه چند پرس‌وجو (MQA) تعداد سرهای KV را کاهش می‌دهند تا در حافظه صرفه‌جویی شود و در طول استنتاج محاسبه انجام شود.

jax.nn.dot_product_attention هر سه مورد را با سرهای KV که به طور خودکار هنگام K < N پخش می‌شوند، مدیریت می‌کند. 

rk = jax.random.key(42)
rk1, rk2, rk3, rk4, rk5 = jax.random.split(rk, 5)

q_mha = jax.random.normal(rk1, (2, 64, 8, 64), dtype=jnp.float32)

# MHA: 8 KV heads
k_mha = jax.random.normal(rk2, (2, 64, 8, 64), dtype=jnp.float32)
v_mha = jax.random.normal(rk3, (2, 64, 8, 64), dtype=jnp.float32)

# GQA: 2 KV heads (each shared by 4 query heads)
k_gqa = jax.random.normal(rk2, (2, 64, 2, 64), dtype=jnp.float32)
v_gqa = jax.random.normal(rk3, (2, 64, 2, 64), dtype=jnp.float32)

# MQA: 1 KV head (shared by all 8 query heads)
k_mqa = jax.random.normal(rk4, (2, 64, 1, 64), dtype=jnp.float32)
v_mqa = jax.random.normal(rk5, (2, 64, 1, 64), dtype=jnp.float32)

out_mha = block_tree(jax.nn.dot_product_attention(q_mha, k_mha, v_mha))
out_gqa = block_tree(jax.nn.dot_product_attention(q_mha, k_gqa, v_gqa))
out_mqa = block_tree(jax.nn.dot_product_attention(q_mha, k_mqa, v_mqa))

show_table(
    ["Pattern", "Q shape", "K shape", "V shape", "Output shape"],
    [
        ("MHA", q_mha.shape, k_mha.shape, v_mha.shape, out_mha.shape),
        ("GQA", q_mha.shape, k_gqa.shape, v_gqa.shape, out_gqa.shape),
        ("MQA", q_mha.shape, k_mqa.shape, v_mqa.shape, out_mqa.shape),
    ],
    title="Multi-head attention variants — all should produce the same output shape",
)

هر سه ردیف باید شکل خروجی یکسانی را گزارش دهند، (2, 64, 8, 64) ، حتی اگر شکل‌های K و V از 8 سر به 2 به 1 کاهش یابند. نکته همین است: شما می‌توانید کش KV را بدون تغییر چیزی در پایین‌دست توجه، حذف کنید.

۱۰. بنچمارک NVIDIA TransformerEngine و FP8

موتور تبدیل‌کننده‌ی انویدیا (TransformerEngine) ماژول‌های توجه ترکیبی (fused attention) را ارائه می‌دهد که برای پردازنده‌های گرافیکی انویدیا بهینه شده‌اند. ادغام JAX از ماژول‌های به سبک کتان (Flax Linen) استفاده می‌کند (نه NNX)، بنابراین ماژول یک بار مقداردهی اولیه می‌شود و سپس با متغیرهایش اعمال می‌شود.

طول توالی را با TransformerEngine جارو کنید

این سلول بررسی می‌کند که آیا NVIDIA TransformerEngine در دسترس است یا خیر، سپس معیار bf16 DotProductAttention سببی آن را در طول چندین توالی در مقایسه با JAX SDPA با بک‌اندهای XLA و cuDNN در حجم کاری یکسان، ارزیابی می‌کند. 

TE_SEQ_LENS = [128, 256, 512, 1024, 2048]
TE_BATCH = BATCH

HAS_TE = False

try:
    import transformer_engine.jax as te
    import transformer_engine.jax.flax as te_flax
    HAS_TE = True
except ImportError:
    print("TransformerEngine not installed — skipping TE sections.")

if HAS_TE:
    te_results = []

    for sl in TE_SEQ_LENS:
        rk = jax.random.key(sl + 1000)
        rk1, rk2, rk3 = jax.random.split(rk, 3)

        q_te = jax.random.normal(
            rk1, (TE_BATCH, sl, NUM_HEADS, HEAD_DIM), dtype=jnp.bfloat16
        )
        k_te = jax.random.normal(
            rk2, (TE_BATCH, sl, NUM_HEADS, HEAD_DIM), dtype=jnp.bfloat16
        )
        v_te = jax.random.normal(
            rk3, (TE_BATCH, sl, NUM_HEADS, HEAD_DIM), dtype=jnp.bfloat16
        )
        q_te, k_te, v_te = jax.device_put((q_te, k_te, v_te), device)

        te_attention = te_flax.DotProductAttention(
            head_dim=HEAD_DIM,
            num_attention_heads=NUM_HEADS,
            num_gqa_groups=NUM_HEADS,
            attn_mask_type="causal",
            transpose_batch_sequence=False,
        )

        te_vars = te_attention.init(
            jax.random.key(0),
            q_te,
            k_te,
            v_te,
            deterministic=True,
        )

        def te_fn(q, k, v):
            return te_attention.apply(te_vars, q, k, v, deterministic=True)

        row = {"seq_len": sl}

        row["sdpa_xla_bf16_ms"] = benchmark_attention(
            lambda q, k, v: jax.nn.dot_product_attention(
                q, k, v, implementation="xla", is_causal=True
            ),
            q_te, k_te, v_te,
            warmup=2,
            repeats=20,
        )

        if HAS_CUDNN_SDPA:
            row["sdpa_cudnn_bf16_ms"] = benchmark_attention(
                lambda q, k, v: jax.nn.dot_product_attention(
                    q, k, v, implementation="cudnn", is_causal=True
                ),
                q_te, k_te, v_te,
                warmup=2,
                repeats=20,
            )

        row["te_bf16_ms"] = benchmark_attention(
            te_fn,
            q_te, k_te, v_te,
            warmup=2,
            repeats=20,
        )

        te_results.append(row)


    headers = ["Seq len", "SDPA XLA bf16 causal (ms)"]
    if HAS_CUDNN_SDPA:
        headers.append("SDPA cuDNN bf16 causal (ms)")
    headers.append("TE DotProductAttention bf16 causal (ms)")

    table_rows = []
    for r in te_results:
        row = [
            r["seq_len"],
            f"{r['sdpa_xla_bf16_ms']:.2f}",
        ]

        if HAS_CUDNN_SDPA:
            row.append(f"{r['sdpa_cudnn_bf16_ms']:.2f}")

        row.append(f"{r['te_bf16_ms']:.2f}")
        table_rows.append(row)

    show_table(
        headers,
        table_rows,
        title=f"TransformerEngine sequence-length sweep — B={TE_BATCH}, N={NUM_HEADS}, H={HEAD_DIM}",
        aligns=["right"] * len(headers),
    )


    fig, ax = plt.subplots(figsize=(8, 5))
    seq_lens = [r["seq_len"] for r in te_results]

    ax.plot(
        seq_lens,
        [r["sdpa_xla_bf16_ms"] for r in te_results],
        "d-",
        label="SDPA XLA bf16 causal",
        color="#8250df",
    )

    if HAS_CUDNN_SDPA:
        ax.plot(
            seq_lens,
            [r["sdpa_cudnn_bf16_ms"] for r in te_results],
            "^-",
            label="SDPA cuDNN bf16 causal",
            color="#1a7f37",
        )

    ax.plot(
        seq_lens,
        [r["te_bf16_ms"] for r in te_results],
        "o-",
        label="TE DotProductAttention bf16 causal",
        color="#d1242f",
    )

    ax.set_xlabel("Sequence length")
    ax.set_ylabel("Median ms per call")
    ax.set_title("Causal attention latency vs sequence length")
    ax.legend()
    ax.grid(True, alpha=0.25)
    ax.set_xticks(seq_lens)

    fig.tight_layout()
    plt.show()

شما باید به ازای هر طول دنباله از ۱۲۸ تا ۲۰۴۸، یک ردیف جدول و یک نقطه نمودار دریافت کنید. این بنچمارک bf16، TransformerEngine را با JAX SDPA در حجم کاری یکسانِ توجه-علّی مقایسه می‌کند، اما پتانسیل کامل عملکرد TransformerEngine معمولاً در پردازنده‌های گرافیکی Hopper و Blackwell زمانی که autocast FP8 در دسترس است، مشاهده می‌شود.

مسیر FP8 را بررسی کنید

در پردازنده‌های گرافیکی Hopper (با قابلیت محاسبه >= 9.0، مانند H100)، TransformerEngine می‌تواند برای افزایش توان عملیاتی، attention را در FP8 اجرا کند. FP8 از دستورالعمل DelayedScaling استفاده می‌کند که تاریخچه حداکثر مطلق هر تنسور را برای محاسبه فاکتورهای مقیاس‌بندی پویا ردیابی می‌کند.

  • قالب E4M3 برای مسیر رو به جلو (۴ بیت توان، ۳ بیت مانتیس)
  • قالب E5M2 برای مسیر معکوس (۵ بیت توان، ۲ بیت مانتیس)

اگر پردازنده‌ی گرافیکی (GPU) از FP8 پشتیبانی نکند، این سلول نشان می‌دهد که کد بدون اجرا چگونه به نظر می‌رسد. 

if HAS_TE:
    from transformer_engine.common.recipe import DelayedScaling, Format

    gpu_name = f"{device} {getattr(device, 'device_kind', '')}".lower()
    HAS_FP8 = any(
        tag in gpu_name
        for tag in ["h100", "h200", "b100", "b200", "gb200", "blackwell"]
    )

    fp8_recipe = DelayedScaling(
        margin=0,
        fp8_format=Format.HYBRID,
        amax_history_len=1024,
        amax_compute_algo="max",
    )

    if HAS_FP8:
        FP8_SEQ_LEN = 2048
        FP8_BATCH = BATCH

        rk = jax.random.key(9000)
        rk1, rk2, rk3 = jax.random.split(rk, 3)

        q_fp8 = jax.random.normal(
            rk1, (FP8_BATCH, FP8_SEQ_LEN, NUM_HEADS, HEAD_DIM), dtype=jnp.bfloat16
        )
        k_fp8 = jax.random.normal(
            rk2, (FP8_BATCH, FP8_SEQ_LEN, NUM_HEADS, HEAD_DIM), dtype=jnp.bfloat16
        )
        v_fp8 = jax.random.normal(
            rk3, (FP8_BATCH, FP8_SEQ_LEN, NUM_HEADS, HEAD_DIM), dtype=jnp.bfloat16
        )
        q_fp8, k_fp8, v_fp8 = jax.device_put((q_fp8, k_fp8, v_fp8), device)

        fp8_attention = te_flax.DotProductAttention(
            head_dim=HEAD_DIM,
            num_attention_heads=NUM_HEADS,
            num_gqa_groups=NUM_HEADS,
            attn_mask_type="causal",
            transpose_batch_sequence=False,
        )

        bf16_vars = fp8_attention.init(
            jax.random.key(0),
            q_fp8,
            k_fp8,
            v_fp8,
            deterministic=True,
        )

        bf16_out = block_tree(
            fp8_attention.apply(
                bf16_vars,
                q_fp8,
                k_fp8,
                v_fp8,
                deterministic=True,
            )
        )

        with te.autocast(enabled=True, recipe=fp8_recipe):
            fp8_vars = fp8_attention.init(
                jax.random.key(1),
                q_fp8,
                k_fp8,
                v_fp8,
                deterministic=True,
            )
            fp8_out = block_tree(
                fp8_attention.apply(
                    fp8_vars,
                    q_fp8,
                    k_fp8,
                    v_fp8,
                    deterministic=True,
                )
            )

        max_diff_fp8 = float(jnp.max(jnp.abs(
            bf16_out.astype(jnp.float32) - fp8_out.astype(jnp.float32)
        )))

        show_table(
            ["", "Value"],
            [
                ("GPU", getattr(device, "device_kind", str(device))),
                ("Input dtype", str(q_fp8.dtype)),
                ("bf16 output dtype", str(bf16_out.dtype)),
                ("FP8 autocast output dtype", str(fp8_out.dtype)),
                ("Output shape", str(fp8_out.shape)),
                ("Max |TE bf16 - TE FP8 autocast|", f"{max_diff_fp8:.2e}"),
            ],
            title="FP8 attention with TransformerEngine",
        )

    else:
        show_table(
            ["", "Value"],
            [
                ("GPU", getattr(device, "device_kind", str(device))),
                ("FP8 support", "No detected support; requires Hopper/Blackwell-class GPU"),
            ],
            title="FP8 attention — not available on this GPU",
        )

        print()
        print("The FP8 path uses TransformerEngine autocast:")
        print()
        print("  with te.autocast(enabled=True, recipe=fp8_recipe):")
        print("      out = fp8_attention.apply(vars, q, k, v, deterministic=True)")

else:
    print("TransformerEngine not available — FP8 section skipped.")

در L4 باید جدولی را ببینید که GPU شما را نامگذاری کرده و عدم پشتیبانی از FP8 شناسایی شده را گزارش می‌دهد، و پس از آن دو خط چاپ شده که فراخوانی te.autocast را که در یک GPU Hopper استفاده می‌کنید، نشان می‌دهند. آن قطعه کد را نگه دارید: این تنها تغییری است که FP8 در محل فراخوانی نیاز دارد.

۱۱. تمیز کردن

حجم کاری Jupyter، شامل LoadBalancer و درایو دائمی را حذف کنید: 

kubectl delete -f deploy/jupyter.yaml

کلاستر، Node Pool، VPC و حساب کاربری سرویس را از بین ببرید: 

cd terraform
terraform destroy

وقتی از شما خواسته شد، yes را تایپ کنید، سپس تأیید کنید که چیزی جا نمانده است: 

gcloud container clusters list
gcloud compute instances list

هر دو باید برای این پروژه خالی باشند. اگر فقط برای این مجموعه پروژه‌ای ایجاد کرده‌اید، می‌توانید کل پروژه را از کنسول Cloud حذف کنید.

۱۲. تبریک

شما از یک محاسبه‌ی توجه دست‌نویس به هسته‌های ترکیبی بهینه‌سازی‌شده برای پردازنده‌ی گرافیکی (GPU) منتقل شدید و تفاوت را روی سخت‌افزار واقعی اندازه‌گیری کردید.

آنچه آموخته‌اید

  • توجه ساده کار می‌کند اما چندین هسته پردازنده گرافیکی را اجرا می‌کند و ماتریس توجه کامل را در حافظه محقق می‌سازد.
  • jax.nn.dot_product_attention محاسبات را در یک عملیات واحد ادغام می‌کند و با implementation=None JAX بهترین backend را به طور خودکار انتخاب می‌کند.
  • implementation="cudnn" هسته‌های توجه ترکیبی cuDNN انویدیا را مجبور می‌کند، که در توالی‌های طولانی سریع‌ترین هستند و به ورودی‌های bfloat16 یا float16 و قابلیت محاسبه ۸.۰ یا جدیدتر نیاز دارند.
  • پوشش سببی با is_causal=True در هسته ادغام‌شده تعبیه شده است - نیازی به ماتریس پوشش دستی نیست
  • GQA و MQA هدهای KV را کاهش می‌دهند تا در طول استنتاج در حافظه صرفه‌جویی شود، و dot_product_attention پخش را به طور خودکار مدیریت می‌کند.
  • TransformerEngine با دقت اختیاری FP8 در پردازنده‌های گرافیکی Hopper، توجه یکپارچه را فراهم می‌کند.

مراحل بعدی

  • Codelab 6: آموزش JAX مقیاس‌پذیر در چندین پردازنده گرافیکی، به شما نشان می‌دهد که چگونه آرایه‌های شاردینگ را در هر دو L4 و مراحل آموزش را به صورت موازی اجرا کنید.
  • HEAD_DIM تغییر دهید (۳۲، ۶۴، ۱۲۸ را امتحان کنید) و ببینید cuDNN کدام مقادیر را می‌پذیرد و زمان‌بندی‌ها چگونه تغییر می‌کنند.
  • SEQ_LEN را افزایش دهید (۲۵۶، ۵۱۲، ۱۰۲۴، ۲۰۴۸ را امتحان کنید) و میزان استفاده از حافظه و نسبت افزایش سرعت cuDNN را مشاهده کنید.
  • مقایسه MHA، GQA و MQA را با K=2 و K=1 KV head در مقابل 8 query head دوباره اجرا کنید و برای استنتاج، دلیل صرفه‌جویی در حافظه نهان KV را بیان کنید.

اسناد مرجع