cuDNN और TransformerEngine की मदद से, जीपीयू पर अटेंशन लेयर की प्रोसेसिंग को तेज़ करना

1. परिचय

जीपीयू पर Jax का लर्निंग पाथ. लैब 5: जीपीयू पर अटेंशन.

इस कोडलैब में, NVIDIA जीपीयू पर अटेंशन की कंप्यूटेशन पर फ़ोकस किया जाता है. सबसे पहले, एक आसान तरीका इस्तेमाल करें. इसके बाद, JAX के पहले से मौजूद फ़्यूज़्ड अटेंशन का इस्तेमाल करें. इसके बाद, उसी कार्रवाई के लिए NVIDIA के cuDNN बैकएंड को चालू करें. इस लेख के आखिर तक आपको पता चल जाएगा कि हर लेवल पर क्या बदलाव होते हैं. साथ ही, यह भी पता चलेगा कि समस्या बढ़ने पर, जीपीयू के लिए ऑप्टिमाइज़ किए गए पाथ कितनी तेज़ी से काम करते हैं.

आपको क्या करना होगा

  • JAX के बुनियादी ऑपरेशंस की मदद से, स्केल किए गए डॉट-प्रॉडक्ट अटेंशन को शुरू से लागू करना
  • इसे jax.nn.dot_product_attention से बदलें. यह JAX का पहले से मौजूद फ़्यूज़्ड कर्नेल है
  • implementation="cudnn" के साथ cuDNN बैकएंड को लागू करें और कॉज़ल मास्किंग जोड़ें
  • स्वीप सीक्वेंस की लंबाई और बैच साइज़, ताकि यह देखा जा सके कि फ़्यूज़्ड कर्नेल कहां काम आते हैं
  • एमएचए, जीक्यूए, और एमक्यूए जैसे मल्टी-हेड अटेंशन शेप की तुलना करना
  • NVIDIA TransformerEngine के अटेंशन की परफ़ॉर्मेंस की जांच करना और उसके FP8 कोड पाथ की जांच करना

आपको किन चीज़ों की ज़रूरत होगी

  • बिलिंग की सुविधा वाला Google Cloud प्रोजेक्ट. साथ ही, वर्कशॉप क्रेडिट या जीपीयू के इस्तेमाल के लिए रिज़र्वेशन
  • आपके चुने गए इलाके में, कम से कम दो NVIDIA L4 जीपीयू के लिए कोटा (जीपीयू कोटा देखने का तरीका)
  • कोडलैब 1 से 4 या CUDA की सुविधा वाले JAX जीपीयू एनवायरमेंट को पूरा करना
  • मैट्रिक्स मल्टीप्लिकेशन और सॉफ़्टमैक्स के बारे में जानकारी. CUDA या cuDNN का अनुभव होना ज़रूरी नहीं है.

नंबर पोर्ट होने में लगने वाला अनुमानित समय: 60 मिनट.

अटेंशन की सुविधा कैसे काम करती है

स्केल्ड डॉट-प्रॉडक्ट अटेंशन में तीन इनपुट होते हैं — क्वेरी, कुंजी, और वैल्यू — और यह कंप्यूट करता है:

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

स्केलिंग फ़ैक्टर 1 / sqrt(d_k), हेड डाइमेंशन बढ़ने पर डॉट प्रॉडक्ट को बहुत ज़्यादा बढ़ने से रोकता है. इससे सॉफ़्टमैक्स को उन क्षेत्रों में धकेल दिया जाता है जहां उसके ग्रेडिएंट छोटे होते हैं.

चार चरण, जिनमें से हर एक अगले चरण से जुड़ा है:

  1. स्कोर — Q @ K.T: डॉट प्रॉडक्ट से यह पता चलता है कि हर क्वेरी पोज़िशन को हर की पोज़िशन पर कितना ध्यान देना चाहिए.
  2. स्केल + सॉफ़्टमैक्स — softmax(scores / sqrt(d)): स्केलिंग से ग्रेडिएंट गायब होने से बचता है. साथ ही, सॉफ़्टमैक्स स्कोर को अटेंशन वेट में बदल देता है. इनका कुल योग 1 होता है.
  3. अटेंड — weights @ V: वैल्यू वेक्टर के वेटेड सम से, हर क्वेरी पोज़िशन के लिए आउटपुट जनरेट होता है.
  4. आउटपुट — इसका शेप Q जैसा ही होता है: हर क्वेरी पोज़िशन में अब उन पोज़िशन की जानकारी होती है जिन पर उसने ध्यान दिया था.

2. शुरू करने से पहले

अपना प्रोजेक्ट चुनें

Google Cloud console में, बिलिंग की सुविधा वाला कोई प्रोजेक्ट चुनें या बनाएं.

Cloud Shell खोलें

Cloud Shell सेशन शुरू करने के लिए, Cloud Shell चालू करें (कंसोल के सबसे ऊपर दाईं ओर मौजूद टर्मिनल आइकॉन) पर क्लिक करें. इसके बाद, इसे अपने प्रोजेक्ट पर ले जाएं:

gcloud config set project <YOUR_PROJECT_ID>

जीपीयू एनवायरमेंट को सेट अप करना

Cloud Shell में यह कमांड चलाएं. Codelab 1 में हर कमांड के बारे में ज़्यादा जानकारी दी गई है. इसमें जीपीयू कोटा और ज़ोन से जुड़ी ज़रूरी शर्तें भी शामिल हैं.

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 = "" सेट करें. इसके बाद, क्लस्टर को प्रोविज़न करें और JupyterLab को डिप्लॉय करें:

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

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

terraform apply से जाने में करीब 12 मिनट लगते हैं. इसके बाद, Pod और LoadBalancer के शुरू होने का इंतज़ार करें. फिर Pod लॉग से, एक बार इस्तेमाल होने वाला JupyterLab टोकन पढ़ें:

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 खोलें, टोकन चिपकाएं, और /workspace में एक नई Python 3 नोटबुक बनाएं. इस कोडलैब का हर कोड ब्लॉक, उस नोटबुक की सेल में जाता है.

इस कोडलैब के लिए ज़रूरी टूल इंस्टॉल करना

!pip install --quiet matplotlib flax

जीपीयू सेट अप करना और उसकी पुष्टि करना

JAX को इंपोर्ट करें और पुष्टि करें कि डिफ़ॉल्ट बैकएंड, जीपीयू है. इस सेल में 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, CUDA डिवाइसों की सूची, और वह जीपीयू दिखना चाहिए जिसका इस्तेमाल बाकी के कोडलैब में किया जाता है.

3. Q, K, और V टेस्ट ऐरे बनाएं

फ़्यूज़्ड अटेंशन कर्नेल इस बात का ध्यान रखते हैं कि इनपुट ऐरे, मेमोरी में किस तरह लेआउट किए गए हैं. इसलिए, कोई भी अटेंशन कोड लिखने से पहले, Q, K, और V को उस लेआउट में बनाएं जिसकी उम्मीद JAX करता है. jax.nn.dot_product_attention को इस लेआउट की उम्मीद है:

हल्की रोशनी

मतलब

हमारा डिफ़ॉल्ट

B

बैच साइज़

4

T

क्वेरी सीक्वेंस की लंबाई

128

S

कुंजी/वैल्यू के क्रम की लंबाई

128 (सेल्फ़-अटेंशन के लिए T के बराबर)

नहीं

अटेंशन हेड की संख्या

8

H

हर हेड का डाइमेंशन

64

आप 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) और dtype float32 दिखेगा. jax.device_put कॉल, तीनों ऐरे को उस जीपीयू पर पिन करता है जिसे आपने सेटअप सेल में चुना है. इसलिए, बाद में बेंचमार्क में कुछ भी होस्ट ट्रांसफ़र को मेज़र नहीं करता.

4. अटेंशन मैकेनिज़्म को नए सिरे से लागू करना

फ़्यूज़्ड कर्नल का इस्तेमाल करने से पहले, फ़ॉर्मूला सीधे लिखें, ताकि आपको पता चल सके कि कर्नल किस चीज़ को बदलता है. यह तरीका, क्रम से इन चार चरणों का पालन करता है: स्कोर की गिनती करना, स्केल करना, सॉफ़्टमैक्स, और वैल्यू से गुणा करना. हेड और सीक्वेंस डाइमेंशन को ट्रांसपोज़ किया जाता है, ताकि 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 के शेप जैसा ही है. फ़ॉर्मूला के आउटपुट चरण में यही बताया गया था.

5. dot_product_attention पर स्विच करें

jax.nn.dot_product_attention, स्कोर, स्केल, सॉफ़्टमैक्स, और अटेंड चरणों को एक ही कार्रवाई में जोड़ता है. इसके बाद, JAX और XLA, मेमोरी ऐक्सेस पैटर्न को ऑप्टिमाइज़ कर सकते हैं. खास तौर पर, जब सीक्वेंस लंबा हो, तब वे पूरे अटेंशन वेट मैट्रिक्स को लागू करने से बच सकते हैं.

डिफ़ॉल्ट implementation=None के साथ, JAX अपने-आप सबसे अच्छा उपलब्ध बैकएंड चुनता है. 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 दिखना चाहिए. साथ ही, ज़्यादा से ज़्यादा एब्सोलूट डिफ़रेंस कम होना चाहिए और atol=1e-3 की क्लोज़नेस जांच के लिए True दिखना चाहिए.

6. cuDNN फ़्यूज़्ड बैकएंड को ज़बरदस्ती लागू करें

JAX को बैकएंड चुनने की अनुमति देना आसान है. हालांकि, बाहर से यह नहीं बताया जा सकता कि असल में कौनसा कर्नेल चला. implementation="cudnn" सेट करने पर, NVIDIA के cuDNN फ़्यूज़्ड अटेंशन कर्नेल को लागू किया जाता है. ये मैन्युअल तरीके से ऑप्टिमाइज़ किए गए जीपीयू कर्नेल होते हैं. ये अटेंशन कंप्यूटेशन को एक ही कर्नेल लॉन्च में फ़्यूज़ करते हैं. साथ ही, मेमोरी ऐक्सेस पैटर्न को ऑप्टिमाइज़ करते हैं.

ध्यान दें कि cuDNN फ़्यूज़्ड अटेंशन के लिए, हार्डवेयर और शेप से जुड़ी ज़रूरी शर्तें होती हैं. जैसे, जीपीयू कंप्यूट की क्षमता >= 8.0 (Ampere या इसके बाद का वर्शन) या इनपुट dtypes के तौर पर float16 या bfloat16 होनी चाहिए. अगर ज़रूरी शर्तें पूरी नहीं होती हैं और आपने 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 को नहीं चला सकता, तो कोड इसकी वजह प्रिंट करता है. इसके बाद, कोडलैब XLA बैकएंड पर जारी रहता है.

7. कॉज़ल मास्किंग जोड़ना

ऑटोरिग्रेसिव मॉडल (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",
)

आपको पहली पोज़िशन पर बड़ा अंतर और आखिरी पोज़िशन पर शून्य के करीब अंतर दिखना चाहिए. यह असमानता, पोज़िशन 0 वाले मास्क की वजह से होती है. यह मास्क, अपनी कुंजी को छोड़कर बाकी सभी कुंजियों को खो देता है. हालांकि, आखिरी पोज़िशन पहले से ही पूरे क्रम को देख सकती है.

8. कुछ बेंचमार्क टेस्ट रन किए जा रहे हैं

अब आपके पास एक ही फ़ंक्शन को कैलकुलेट करने के तीन तरीके हैं. कौनसा तरीका अपनाना है, यह पूरी तरह से समस्या पर निर्भर करता है. इसका आकलन किया जा सकता है.

डिफ़ॉल्ट शेप में वैरिएंट का समय तय करना

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) अटेंशन मैट्रिक्स को जीपीयू मेमोरी में सेव करने के लिए, ज़्यादा मेमोरी की ज़रूरत होती है. यह मेमोरी O(T²) होती है. cuDNN FlashAttention जैसे फ़्यूज़्ड कर्नेल, कंप्यूटेशन को टाइल करते हैं. इसलिए, वे कभी भी पूरी मैट्रिक्स को मैटीरियलाइज़ नहीं करते. इससे मेमोरी O(T) पर बनी रहती है.

नीचे दिए गए स्वीप में, 64 से 1024 तक के सीक्वेंस की लंबाई के लिए, हर इंप्लिमेंटेशन का समय दिखाया गया है.

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),
)

स्वीप में कुछ समय लगता है, क्योंकि हर सीक्वेंस की लंबाई, समय तय करने से पहले तीन या चार अलग-अलग तरीकों को कंपाइल करती है. आपको 64 से 1024 तक की हर सीक्वेंस लंबाई के लिए, टेबल की एक पंक्ति मिलनी चाहिए.

इस प्लॉट में, सीक्वेंस की लंबाई के साथ अटेंशन लेटेंसी में होने वाले बदलावों की तुलना की गई है. इसमें, float32 और bfloat16 में JAX/XLA SDPA, और bfloat16 में cuDNN फ़्यूज़्ड अटेंशन के लिए, सामान्य तरीके से लागू किए गए बदलावों की तुलना की गई है.

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()

आपको हर लागू किए गए तरीके के लिए एक लाइन दिखेगी. आपको दिखेगा कि जैसे-जैसे क्रम लंबा होता जाता है, वैसे-वैसे गैप बढ़ता जाता है.

बैच साइज़ को स्वीप करना

बड़े बैच, कर्नेल लॉन्च ओवरहेड को कम करते हैं और जीपीयू के इस्तेमाल को बेहतर बनाते हैं. हालांकि, ऐसा तब तक होता है, जब तक जीपीयू मेमोरी की कमी न हो जाए. नीचे दिए गए स्वीप में, सीक्वेंस की लंबाई 256 पर तय की गई है. साथ ही, बैच का साइज़ अलग-अलग है.

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()

आपको एक टेबल और एक प्लॉट मिलना चाहिए. इसमें 1 से 16 तक के हर बैच साइज़ के लिए एक एंट्री होनी चाहिए. इस स्वीप में, सामान्य तरीके से लागू किए गए मॉडल को शामिल नहीं किया गया है. इसलिए, दोनों आउटपुट सिर्फ़ तीन फ़्यूज़ किए गए पाथ की तुलना करते हैं.

9. एमएचए, जीक्यूए, और एमक्यूए की तुलना करना

अब तक हर हेडर में अपनी कुंजियां और वैल्यू होती थीं. हालांकि, डिकोडर अक्सर इस सिमिट्री को तोड़ देते हैं. ऐसा इसलिए होता है, क्योंकि अनुमान लगाने के दौरान, की-वैल्यू कैश मेमोरी का इस्तेमाल सबसे ज़्यादा होता है.

मल्टी-हेड अटेंशन (एमएचए), हर हेड को अपना Q, K, और V प्रोजेक्शन देता है. ग्रुप-क्वेरी अटेंशन (GQA) और मल्टी-क्वेरी अटेंशन (MQA), केवी हेड की संख्या को कम करते हैं. इससे मेमोरी की बचत होती है और अनुमान लगाने के दौरान कंप्यूटेशन की ज़रूरत कम होती है.

jax.nn.dot_product_attention, केवी हेड के साथ तीनों को मैनेज करता है. 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 शेप आठ हेड से घटकर दो और फिर एक हो जाएं. यही तो बात है: अटेंशन के बाद के किसी भी हिस्से में बदलाव किए बिना, केवी कैश को कम किया जा सकता है.

10. NVIDIA TransformerEngine और FP8 के लिए मानदंड

NVIDIA का TransformerEngine, NVIDIA के जीपीयू के लिए ऑप्टिमाइज़ किए गए फ़्यूज़्ड अटेंशन मॉड्यूल उपलब्ध कराता है. JAX इंटिग्रेशन, Flax Linen-स्टाइल वाले मॉड्यूल (NNX नहीं) का इस्तेमाल करता है. इसलिए, मॉड्यूल को एक बार शुरू किया जाता है और फिर उसके वैरिएबल के साथ लागू किया जाता है.

TransformerEngine की मदद से, सीक्वेंस की लंबाई को स्वीप करना

यह सेल जांच करती है कि NVIDIA TransformerEngine उपलब्ध है या नहीं. इसके बाद, यह एक ही वर्कलोड पर XLA और cuDNN बैकएंड के साथ JAX SDPA के मुकाबले, कई सीक्वेंस लेंथ में इसके bf16 कॉज़ल DotProductAttention को बेंचमार्क करती है.

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()

आपको 128 से 2048 तक की हर सीक्वेंस लेंथ के लिए, एक टेबल लाइन और एक प्लॉट पॉइंट मिलना चाहिए. यह bf16 मानदंड, TransformerEngine की तुलना JAX SDPA से करता है. यह तुलना, एक ही कॉज़ल-अटेंशन वर्कलोड पर की जाती है. हालांकि, TransformerEngine की पूरी परफ़ॉर्मेंस क्षमता आम तौर पर Hopper और Blackwell जीपीयू पर दिखती है. ऐसा तब होता है, जब FP8 ऑटोकास्ट उपलब्ध हो.

FP8 पाथ की जांच करना

हॉपर जीपीयू (कंप्यूट क्षमता >= 9.0, जैसे कि H100) पर, TransformerEngine ज़्यादा थ्रूपुट के लिए FP8 में अटेंशन चला सकता है. FP8, DelayedScaling रेसिपी का इस्तेमाल करता है. यह रेसिपी, डाइनैमिक स्केलिंग फ़ैक्टर का हिसाब लगाने के लिए, हर टेंसर के ऐब्सलूट-मैक्स इतिहास को ट्रैक करती है.

  • फ़ॉरवर्ड पास के लिए E4M3 फ़ॉर्मैट (4 एक्सपोनेंट, 3 मैंटीसा बिट)
  • बैकवर्ड पास के लिए E5M2 फ़ॉर्मैट (5 एक्सपोनेंट, 2 मैंटीसा बिट)

अगर जीपीयू, 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 पर, आपको एक टेबल दिखेगी. इसमें आपके जीपीयू का नाम होगा और यह जानकारी होगी कि FP8 के लिए कोई सपोर्ट नहीं मिला है. इसके बाद, दो प्रिंट की गई लाइनें दिखेंगी. इनमें वह te.autocast कॉल दिखेगा जिसका इस्तेमाल Hopper जीपीयू पर किया जाएगा. उस स्निपेट को बनाए रखें: कॉल साइट पर FP8 के लिए सिर्फ़ यही बदलाव ज़रूरी है.

11. व्यवस्थित करें

Jupyter वर्कलोड को मिटाएं. इसमें LoadBalancer और पर्सिस्टेंट वॉल्यूम शामिल है:

kubectl delete -f deploy/jupyter.yaml

क्लस्टर, नोड पूल, वीपीसी, और सेवा खाते को मिटाएं:

cd terraform
terraform destroy

जब कहा जाए, तब yes टाइप करें. इसके बाद, पुष्टि करें कि कोई भी आइटम नहीं छूटा है:

gcloud container clusters list
gcloud compute instances list

इस प्रोजेक्ट के लिए, दोनों फ़ील्ड खाली होने चाहिए. अगर आपने सिर्फ़ इस सीरीज़ के लिए कोई प्रोजेक्ट बनाया है, तो Cloud console में जाकर पूरे प्रोजेक्ट को मिटाया जा सकता है.

12. बधाई हो

आपने हाथ से लिखे गए अटेंशन कंप्यूटेशन से, जीपीयू के लिए ऑप्टिमाइज़ किए गए फ़्यूज़्ड कर्नेल पर स्विच किया और असल हार्डवेयर पर अंतर को मेज़र किया.

आपने क्या सीखा

  • नेव अटेंशन काम करता है, लेकिन यह कई जीपीयू कर्नेल लॉन्च करता है और मेमोरी में पूरा अटेंशन मैट्रिक्स बनाता है
  • jax.nn.dot_product_attention, कंप्यूटेशन को एक ही ऑपरेशन में फ़्यूज़ करता है. साथ ही, implementation=None JAX अपने-आप सबसे सही बैकएंड चुनता है
  • implementation="cudnn", NVIDIA के cuDNN फ़्यूज़्ड अटेंशन कर्नेल को लागू करता है. ये लंबी सीक्वेंस पर सबसे तेज़ होते हैं. इसके लिए, bfloat16 या float16 इनपुट और 8.0 या इसके बाद के वर्शन की कंप्यूट क्षमता ज़रूरी है
  • is_causal=True के साथ कॉज़ल मास्किंग की सुविधा, फ़्यूज़्ड कर्नल में पहले से मौजूद होती है. इसके लिए, मैन्युअल तरीके से मास्क मैट्रिक्स की ज़रूरत नहीं होती
  • GQA और MQA, अनुमान लगाने के दौरान मेमोरी बचाने के लिए केवी हेड को कम करते हैं. साथ ही, dot_product_attention ब्रॉडकास्ट को अपने-आप मैनेज करता है
  • TransformerEngine, Hopper जीपीयू पर फ़्यूज़्ड अटेंशन की सुविधा देता है. इसमें FP8 प्रिसिज़न का विकल्प भी होता है

अगले चरण

  • कोडलैब 6: JAX ट्रेनिंग को कई जीपीयू पर स्केल करना में, आपको दोनों L4 पर ऐरे को शार्ड करने और ट्रेनिंग के चरणों को एक साथ चलाने का तरीका बताया जाएगा
  • HEAD_DIM को बदलें (32, 64, 128 आज़माएं) और देखें कि cuDNN कौनसी वैल्यू स्वीकार करता है और समय कैसे बदलता है
  • SEQ_LEN को बढ़ाएं (256, 512, 1024, 2048 आज़माएं) और मेमोरी के इस्तेमाल और cuDNN की स्पीड बढ़ाने का अनुपात देखें
  • आठ क्वेरी हेड के लिए, K=2 और K=1 केवी हेड के साथ एमएचए, जीक्यूए, और एमक्यूए की तुलना फिर से करें. साथ ही, अनुमान के लिए केवी कैश मेमोरी की बचत के बारे में बताएं

रेफ़रंस के लिए दस्तावेज़