cuDNN ও TransformerEngine-এর সাহায্যে GPU-তে অ্যাটেনশন স্পিড আপ করুন

১. ভূমিকা

Jax on GPU লার্নিং পাথ। ল্যাব ৫: GPU-তে মনোযোগ।

এই কোডল্যাবে আপনি NVIDIA GPU-তে অ্যাটেনশনের কম্পিউটেশন-এর উপর মনোযোগ দেবেন। আপনি একটি সাধারণ ইমপ্লিমেন্টেশন দিয়ে শুরু করেন, এটিকে JAX-এর বিল্ট-ইন ফিউজড অ্যাটেনশন দিয়ে প্রতিস্থাপন করেন, তারপর একই অপারেশনের জন্য NVIDIA-এর cuDNN ব্যাকএন্ড সক্রিয় করেন। শেষে আপনি জানতে পারবেন যে প্রতিটি লেভেলে কী পরিবর্তন হয় এবং সমস্যা বাড়ার সাথে সাথে GPU-অপ্টিমাইজ করা পাথ কতটা দ্রুত হয়।

আপনি কী করবেন

  • প্রাথমিক JAX অপারেশনের মাধ্যমে স্ক্র্যাচ থেকে স্কেল করা ডট-প্রোডাক্ট অ্যাটেনশন প্রয়োগ করা
  • এটি jax.nn.dot_product_attention, JAX-এর বিল্ট-ইন ফিউজড কার্নেল দিয়ে পরিবর্তন করুন
  • implementation="cudnn"-এর সাথে cuDNN ব্যাকএন্ড ফোর্স এবং কজাল মাস্কিং যোগ করো
  • ফিউজড কার্নেল কোথায় কাজে লাগে তা দেখতে সুইপ সিকোয়েন্সের দৈর্ঘ্য ও ব্যাচের সাইজ দেখুন
  • মাল্টি-হেড অ্যাটেনশন শেপ MHA, GQA এবং MQA-এর তুলনা করো
  • NVIDIA TransformerEngine অ্যাটেনশন বেঞ্চমার্ক করো এবং এর FP8 কোড পাথ পরীক্ষা করো

আপনার কী প্রয়োজন হবে

  • বিলিং চালু করা আছে এমন একটি Google Cloud প্রোজেক্ট এবং ওয়ার্কশপ ক্রেডিট বা রিজার্ভেশন যা GPU ব্যবহার কভার করে
  • আপনার বেছে নেওয়া অঞ্চলে কমপক্ষে ২টি NVIDIA L4 GPU-এর জন্য কোটা (GPU কোটা কীভাবে চেক করবেন)
  • ১ থেকে ৪ নম্বর কোডল্যাব সম্পূর্ণ করা অথবা সমতুল্য CUDA-চালিত JAX GPU এনভায়রনমেন্ট
  • ম্যাট্রিক্স মাল্টিপ্লিকেশন ও সফটম্যাক্স সম্পর্কে জানা। CUDA বা cuDNN-এর অভিজ্ঞতা থাকার প্রয়োজন নেই।

সম্পূর্ণ হতে আনুমানিক যতক্ষণ লাগবে: ৬০ মিনিট।

অ্যাটেনশন কীভাবে কাজ করে

স্কেল করা ডট-প্রোডাক্ট অ্যাটেনশন তিনটি ইনপুট নেয় — কোয়েরি, কী এবং ভ্যালু — এবং গণনা করে:

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

স্কেলিং ফ্যাক্টর 1 / sqrt(d_k) ডট প্রোডাক্টকে খুব বেশি বড় হতে দেয় না, কারণ হেড ডাইমেনশন বাড়লে সফটম্যাক্সকে এমন অঞ্চলে ঠেলে দেয় যেখানে এর গ্রেডিয়েন্টগুলি খুব ছোট হয়।

চারটি ধাপ, প্রতিটি ধাপ পরবর্তী ধাপের সাথে যুক্ত:

  1. স্কোর — Q @ K.T: ডট প্রোডাক্ট পরিমাপ করে যে প্রতিটি কোয়েরি পজিশন প্রতিটি কী পজিশনে কতটা মনোযোগ দেবে।
  2. Scale + Softmax — softmax(scores / sqrt(d)): স্কেলিং গ্রেডিয়েন্ট ভ্যানিশিং প্রতিরোধ করে এবং সফটম্যাক্স স্কোরকে অ্যাটেনশন ওয়েটে কনভার্ট করে যার যোগফল ১ হয়।
  3. অ্যাটেন্ড — weights @ V: ভ্যালু ভেক্টরের ওয়েটেড সাম প্রতিটি কোয়েরি পজিশনের জন্য আউটপুট তৈরি করে।
  4. আউটপুট — Q-এর মতো একই আকার: প্রতিটি কোয়েরি পজিশন এখন যে পজিশনে অ্যাটেন্ড করেছে সেখান থেকে তথ্য বহন করে।

২. শুরু করার আগে

আপনার প্রোজেক্ট বেছে নিন

Google Cloud কনসোল থেকে, বিলিং চালু আছে এমন কোনও প্রোজেক্ট বেছে নিন বা তৈরি করুন।

Cloud Shell খুলুন

Cloud Shell সেশন শুরু করতে 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 = "" সেট করুন, তারপর ক্লাস্টার প্রোভিশন করে JupyterLab ডিপ্লয় করুন:

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

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

terraform apply-এ যেতে প্রায় ১২ মিনিট সময় লাগে। এটি শেষ হলে, 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

GPU সেট-আপ ও যাচাই করা

JAX ইমপোর্ট করুন এবং কনফার্ম করুন যে ডিফল্ট ব্যাকএন্ড হল 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, CUDA ডিভাইসের তালিকা এবং কোডল্যাবের বাকি অংশে ব্যবহৃত GPU দেখতে পাবেন।

৩. Q, K এবং V টেস্ট অ্যারে তৈরি করুন

ফিউজড অ্যাটেনশন কার্নেল মেমরিতে ইনপুট অ্যারে কীভাবে সাজানো আছে তা বিবেচনা করে, তাই কোনও অ্যাটেনশন কোড লেখার আগে JAX-এর প্রত্যাশিত লেআউটে Q, K এবং V তৈরি করুন। jax.nn.dot_product_attention নিম্নলিখিত লেআউট আশা করে:

কম আলো

অর্থ

আমাদের ডিফল্ট

B

ব্যাচের সাইজ

4

ত

কোয়েরি সিকোয়েন্সের দৈর্ঘ্য

128

দঃ

কী/ভ্যালু সিকোয়েন্সের দৈর্ঘ্য

১২৮ (সেল্ফ-অ্যাটেনশনের জন্য T-এর মতো একই)

না

অ্যাটেনশন হেডের সংখ্যা

8

H

মাথাপিছু ডাইমেনশন

৬৪

আপনি 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 কল সেট-আপ সেলে আপনার বেছে নেওয়া GPU-তে তিনটি অ্যারে পিন করে, তাই পরে বেঞ্চমার্কে কোনও কিছুই হোস্ট ট্রান্সফার পরিমাপ করে না।

৪. শুরু থেকে মনোযোগ প্রয়োগ করা

ফিউজড কার্নেল ব্যবহার করার আগে, সরাসরি ফর্মুলাটি লিখুন যাতে আপনি দেখতে পারেন যে কার্নেলটি ঠিক কী প্রতিস্থাপন করে। এই প্রয়োগটি ক্রমানুসারে চারটি ধাপ অনুসরণ করে: স্কোর গণনা, স্কেল, সফটম্যাক্স, মান দ্বারা গুণ। আপনি হেড এবং সিকোয়েন্স ডাইমেনশন ট্রান্সপোজ করেন যাতে matmul সিকোয়েন্স অক্ষের উপর কাজ করে।

এটি সঠিকভাবে কাজ করে, তবে এটি স্কোর ম্যাটমুল, সফটম্যাক্স, ভ্যালু ম্যাটমুলের জন্য তিনটি আলাদা GPU কার্নেল লঞ্চ করে এবং এটি GPU মেমরিতে সম্পূর্ণ (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 স্কোর, স্কেল, সফটম্যাক্স এবং অ্যাটেন্ড ধাপগুলিকে একটি অপারেশনে একত্রিত করে। JAX এবং XLA তখন মেমরি অ্যাক্সেস প্যাটার্ন অপ্টিমাইজ করতে পারে। বিশেষ করে, সিকোয়েন্স দীর্ঘ হলে, তারা সম্পূর্ণ অ্যাটেনশন ওয়েট ম্যাট্রিক্স মেটেরিয়ালাইজ করা এড়াতে পারে।

ডিফল্ট implementation=None-এর সাথে, JAX অটোমেটিক সেরা উপলভ্য ব্যাকএন্ড বেছে নেয়। cuDNN উপলভ্য এবং মানানসই ইনপুট সহ GPU-তে, এটি আগে থেকেই 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 ভার্সনে, একটি ছোট সর্বোচ্চ পরম পার্থক্য এবং atol=1e-3 ক্লোজনেস চেকের জন্য True।

৬. cuDNN ফিউজড ব্যাকএন্ড জোর করে চালু করুন

JAX-কে ব্যাকএন্ড বেছে নিতে দেওয়া সুবিধাজনক, কিন্তু বাইরে থেকে আপনি বলতে পারবেন না যে আসলে কোন কার্নেল রান করেছে। implementation="cudnn" সেট করলে NVIDIA-এর cuDNN ফিউজড অ্যাটেনশন কার্নেল ব্যবহার করতে বাধ্য করে। এগুলি হল হ্যান্ড-অপ্টিমাইজ করা GPU কার্নেল যা অপ্টিমাইজ করা মেমরি অ্যাক্সেস প্যাটার্নের সাথে একটি একক কার্নেল লঞ্চে সম্পূর্ণ মনোযোগ গণনাকে ফিউজ করে।

মনে রাখবেন যে cuDNN ফিউজড অ্যাটেনশনের হার্ডওয়্যার ও শেপ সংক্রান্ত প্রয়োজনীয়তা আছে, যেমন GPU কম্পিউট ক্ষমতা >= 8.0 (Ampere বা তার পরবর্তী যেকোনও ভার্সন) বা ইনপুট dtype হিসেবে 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.")

আপনার GPU যদি cuDNN SDPA রান করতে না পারে, তাহলে কোডটি কারণ প্রিন্ট করে এবং কোডল্যাব 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 কিছুটা দ্রুততম হওয়া উচিত, তবে ফিউজড অ্যাটেনশনের সুবিধা সাধারণত দীর্ঘ সিকোয়েন্স দৈর্ঘ্যে আরও স্পষ্ট হয়ে ওঠে, যেখানে সম্পূর্ণ অ্যাটেনশন ম্যাট্রিক্স এড়ানো গুরুত্বপূর্ণ।

সিকোয়েন্সের দৈর্ঘ্য সুইপ করুন

ফিউজড অ্যাটেনশন কার্নেলের সুবিধা সিকোয়েন্সের দৈর্ঘ্যের সাথে সাথে বৃদ্ধি পায়। নাইভ ইমপ্লিমেন্টেশন GPU মেমরিতে (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),
)

সুইপ করতে কিছুটা সময় লাগে, কারণ প্রতিটি সিকোয়েন্সের দৈর্ঘ্য টাইম করার আগে তিন বা চারটি আলাদা ইমপ্লিমেন্টেশন কম্পাইল করে। আপনাকে ৬৪ থেকে ১০২৪ পর্যন্ত প্রতিটি সিকোয়েন্স দৈর্ঘ্যের জন্য একটি করে টেবিল রো তৈরি করতে হবে।

এই প্লটটি তুলনা করে যে কীভাবে অ্যাটেনশন লেটেন্সি সিকোয়েন্সের দৈর্ঘ্যের সাথে পরিবর্তিত হয়, 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()

প্রতিটি ইমপ্লিমেন্টেশনের জন্য একটি করে লাইন দেখতে পাবেন। এবং আপনি দেখতে পাবেন যে সিকোয়েন্সটি দীর্ঘ হওয়ার সাথে সাথে ব্যবধানটি বৃদ্ধি পাচ্ছে।

ব্যাচের সাইজ সুইপ করুন

বড় ব্যাচ কার্নেল লঞ্চ ওভারহেডকে অ্যামর্টাইজ করে এবং GPU মেমরি বাধা হয়ে দাঁড়ানো পর্যন্ত GPU-এর ব্যবহার উন্নত করে। নিচের সুইপ সিকোয়েন্সের দৈর্ঘ্য ২৫৬-এ স্থির রাখে এবং ব্যাচের সাইজ পরিবর্তন করে।

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 শেপ ৮টি হেড থেকে ২ থেকে ১-এ সঙ্কুচিত হয়। এটাই হল আসল কথা: আপনি অ্যাটেনশনের ডাউনস্ট্রিমে কোনও পরিবর্তন না করেই KV ক্যাশে কাট করতে পারবেন।

১০. NVIDIA TransformerEngine এবং FP8 বেঞ্চমার্ক করা

NVIDIA-এর TransformerEngine NVIDIA GPU-এর জন্য অপ্টিমাইজ করা ফিউজড অ্যাটেনশন মডিউল প্রদান করে। 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()

১২৮ থেকে ২০৪৮ পর্যন্ত প্রতিটি সিকোয়েন্স দৈর্ঘ্যের জন্য আপনার একটি টেবিল রো এবং একটি প্লট পয়েন্ট পাওয়া উচিত। এই bf16 বেঞ্চমার্ক একই কজাল-অ্যাটেনশন ওয়ার্কলোডে JAX SDPA-এর সাথে TransformerEngine-এর তুলনা করে, তবে TransformerEngine-এর সম্পূর্ণ পারফর্ম্যান্সের সম্ভাবনা সাধারণত Hopper ও Blackwell GPU-তে দেখা যায় যখন FP8 অটো-কাস্ট উপলভ্য থাকে।

FP8 পাথ চেক করুন

Hopper GPU-তে (কম্পিউট ক্ষমতা >= ৯.০, যেমন H100), TransformerEngine অতিরিক্ত থ্রুপুটের জন্য 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 সাপোর্ট শনাক্ত করা যায়নি বলে রিপোর্ট করা হয়েছে, এর পরে দুটি প্রিন্ট করা লাইন রয়েছে যা Hopper GPU-তে আপনি যে te.autocast কল ব্যবহার করবেন তা দেখায়। সেই স্নিপেট রেখে দিন: কল সাইটে FP8-এর জন্য এটিই একমাত্র পরিবর্তন।

১১. অপ্রয়োজনীয়গুলি মুছুন

LoadBalancer এবং পারসিস্টেন্ট ভলিউম সহ Jupyter ওয়ার্কলোড মুছে দিন:

kubectl delete -f deploy/jupyter.yaml

ক্লাস্টার, নোড পুল, VPC এবং পরিষেবা অ্যাকাউন্ট ধ্বংস করুন:

cd terraform
terraform destroy

প্রম্পট করা হলে yes টাইপ করুন, তারপর কনফার্ম করুন যে কিছু ফেলে যাননি:

gcloud container clusters list
gcloud compute instances list

এই প্রোজেক্টের জন্য দুটিই খালি থাকতে হবে। আপনি যদি শুধু এই সিরিজের জন্য কোনও প্রোজেক্ট তৈরি করে থাকেন, তাহলে আপনি Cloud কনসোল থেকে পুরো প্রোজেক্টটি মুছে ফেলতে পারেন।

১২. অভিনন্দন

আপনি হাতে লেখা অ্যাটেনশন কম্পিউটেশন থেকে GPU-অপ্টিমাইজ করা ফিউজড কার্নেলে চলে গেছেন এবং আসল হার্ডওয়্যারে পার্থক্য পরিমাপ করেছেন।

আপনি যা শিখেছেন

  • নাইভ অ্যাটেনশন কাজ করে তবে একাধিক GPU কার্নেল লঞ্চ করে এবং মেমরিতে সম্পূর্ণ অ্যাটেনশন ম্যাট্রিক্সকে মেটেরিয়ালাইজ করে
  • jax.nn.dot_product_attention কম্পিউটেশনকে একটি অপারেশনে ফিউজ করে এবং implementation=None JAX-এর মাধ্যমে অটোমেটিক সেরা ব্যাকএন্ড বেছে নেয়
  • implementation="cudnn" NVIDIA-এর cuDNN ফিউজড অ্যাটেনশন কার্নেলকে বাধ্য করে, যা দীর্ঘ সিকোয়েন্সে সবচেয়ে দ্রুত কাজ করে এবং এর জন্য bfloat16 বা float16 ইনপুট ও কম্পিউট ক্ষমতা 8.0 বা তার পরবর্তী ভার্সন প্রয়োজন
  • is_causal=True-এর সাথে কজাল মাস্কিং ফিউজড কার্নেলে তৈরি করা হয়েছে — কোনও ম্যানুয়াল মাস্ক ম্যাট্রিক্সের প্রয়োজন নেই
  • অনুমান করার সময় মেমরি বাঁচাতে GQA এবং MQA KV হেড কমায় এবং dot_product_attention অটোমেটিক ব্রডকাস্টিং ম্যানেজ করে
  • TransformerEngine Hopper GPU-তে ঐচ্ছিক FP8 প্রিসিশন সহ ফিউজড অ্যাটেনশন প্রদান করে

পরবর্তী ধাপ

  • কোডল্যাব ৬: একাধিক GPU জুড়ে JAX ট্রেনিং স্কেল করা আপনাকে L4 জুড়ে অ্যারে শার্ডিং এবং সমান্তরালভাবে ট্রেনিং স্টেপ রান করা দেখাবে
  • HEAD_DIM পরিবর্তন করুন (৩২, ৬৪, ১২৮ ব্যবহার করে দেখুন) এবং cuDNN কোন ভ্যালু গ্রহণ করে ও টাইমিং কীভাবে পরিবর্তিত হয় তা দেখুন
  • SEQ_LEN বাড়ান (২৫৬, ৫১২, ১০২৪, ২০৪৮ ব্যবহার করে দেখুন) এবং মেমরির ব্যবহার ও cuDNN স্পিডআপ অনুপাত দেখুন
  • ৮টি কোয়েরি হেডের বিপরীতে K=2 এবং K=1 KV হেড সহ MHA, GQA এবং MQA তুলনা আবার চালান এবং ইনফারেন্সের জন্য KV ক্যাশে সেভিংস সম্পর্কে কারণ ব্যাখ্যা করুন

রেফারেন্স ডকুমেন্ট