দক্ষ attention — KV cache, FlashAttention-এর ধারণা আর long context

এতদিন আমরা attention-এর অর্থ নিয়ে ভেবেছি: পর্ব ৩২ (Attention) আর পর্ব ৩৩ (Self-attention আর multi-head attention). পর্ব ৩৬-এ (নিজে Transformer লেখা) একটি ছোট model দাঁড়িয়েছে। এই পর্বের প্রশ্ন ভিন্ন: সেই model চালাতে খরচ কত, আর কমাবেন কীভাবে? LLM-এর ধারণা নতুন লাগলে আগে দেখে নিন AI-এর ভিত্তি কোর্সের পর্ব ৭০। পাতার চারটি ব্লক ব্রাউজারের Python আর NumPy-তে চলে; প্রথমবার প্যাকেজ নামতে কয়েক সেকেন্ড লাগে। PyTorch-এর সব আউটপুট ডেস্কটপের সত্যিকারের run থেকে নেওয়া।

১ · একটি নতুন token-এর দাম

LLM লেখে এক token করে। প্রতিটি নতুন token আগের সব token দেখে তৈরি হয়; একে বলে autoregressive generation. সরল কোডটি এমন: এখন পর্যন্ত যত token আছে, সব model-এ দিন, শেষ ঘরের logits থেকে পরের token বাছুন, জুড়ে দিন, আবার।

পর্ব ৩৩-এর হিসাবে T-টি token-এর একটি self-attention-এ শুধু score বানাতেই T × T × d_k × h = T²·d-টি গুণ লাগে। এখানে T token সংখ্যা, h head সংখ্যা, d_k = d/h এক head-এর মাত্রা, আর d পুরো মাত্রা। এই কাজ প্রতিটি নতুন token-এ, প্রতিটি স্তরে হয়। T-টি token লিখতে মোট:

d · (1² + 2² + … + T²) ≈ d · T³ / 3

T দ্বিগুণ হলে কাজ আট গুণ। অথচ একটি নতুন token-এর জন্য দরকার মাত্র একটি সারি: নতুন query-র সঙ্গে সব key-র মিল, T·d-টি গুণ। বাকি সারিগুলো আগেই হিসাব করা, আর ফেলে দেওয়া।

২ · KV cache — কী রাখা হয়, কী হয় না

কেন পুরনো হিসাব আবার কাজে লাগে? কারণটি causal mask. Token ৩ কখনো token ৪-কে দেখে না। তাই token ৪ আসার পরেও token ৩-এর প্রতিটি স্তরের hidden vector হুবহু আগের মতো থাকে, ফলে তার k আর v-ও। পুরনো অংশ গাণিতিকভাবে অপরিবর্তিত, এটি কোনো অনুমান নয়। Mask না থাকলে এই কৌশল চলত না। নিচে একটি ছোট decoder: L = ৪, h = ৪, d = ১২৮, তাই d_k = ৩২. গুণের হিসাব এক স্তরের, সব head মিলিয়ে।

নিজে দেখুন। নিচের ছোট decoder একই seed-এ দুইভাবে লেখে, আর mm ফাংশন প্রতিটি matmul-এর গুণ গোনে। Sampling চালু আছে, যাতে “মিল” কথাটি অর্থপূর্ণ হয়: একটি probability সামান্য বদলালেও sequence আলাদা হতে পারত।

import numpy as np, matplotlib.pyplot as plt
rng = np.random.default_rng(37)
V, d, L = 20, 16, 2                     # vocab, মাত্রা, স্তর; সরলতার জন্য এক head
E, P = rng.normal(0, .5, (V, d)), rng.normal(0, .5, (256, d))
W = rng.normal(0, .3, (L, 3, d, d))     # প্রতি স্তরে Wq, Wk, Wv
count = 0
def mm(A, B):                           # গুণ গুনে রাখা matmul
    global count; count += A.shape[0] * A.shape[1] * B.shape[1]
    return A @ B
def softmax(s):
    e = np.exp(s - s.max(-1, keepdims=True)); return e / e.sum(-1, keepdims=True)
def step(ids, start, cache):            # ids শুধু নতুন token; cache[i] = (K, V)
    x, n = E[ids] + P[start:start + len(ids)], len(ids)
    for i in range(L):
        q, k, v = mm(x, W[i, 0]), mm(x, W[i, 1]), mm(x, W[i, 2])
        if cache[i]: k, v = np.r_[cache[i][0], k], np.r_[cache[i][1], v]
        cache[i] = (k, v)
        s = mm(q, k.T) / np.sqrt(d)
        s[np.triu(np.ones((n, len(k)), bool), len(k) - n + 1)] = -np.inf
        x = x + mm(softmax(s), v)
    return mm(x, E.T)
def generate(prompt, n, use_cache):
    global count; count = 0; g = np.random.default_rng(1)
    ids, new, cache, cost = list(prompt), list(prompt), [None] * L, []
    for _ in range(n):
        lg = (step(np.array(new), len(ids) - len(new), cache) if use_cache
              else step(np.array(ids), 0, [None] * L))
        ids.append(int(g.choice(V, p=softmax(lg[-1] / 8)))); new = ids[-1:]; cost.append(count)
    return ids, np.array(cost)
a, ca = generate([3, 1, 4], 200, False); b, cb = generate([3, 1, 4], 200, True)
print("identical tokens:", a == b, "| first 12:", a[:12])
for n in (50, 100, 200): print(f"T+{n:3d}  no cache {ca[n-1]:>11,}  cache {cb[n-1]:>9,}  ratio {ca[n-1]/cb[n-1]:5.1f}")
print("100 to 200 tokens: no cache x", round(ca[199]/ca[99], 2), "| cache x", round(cb[199]/cb[99], 2))
plt.figure(figsize=(5, 3)); plt.semilogy(ca, label="no cache"); plt.semilogy(cb, label="KV cache")
plt.xlabel("new tokens generated"); plt.ylabel("multiplications (total)"); plt.legend(); plt.show()

identical tokens: True | first 12: [3, 1, 4, 10, 18, 3, 18, 6, 8, 16, 8, 10]
T+ 50  no cache   5,638,400  cache   184,896  ratio  30.5
T+100  no cache  32,716,800  cache   525,696  ratio  62.2
T+200  no cache 215,193,600  cache 1,687,296  ratio 127.5
100 to 200 tokens: no cache x 6.58 | cache x 3.21

৩ · PyTorch-এ সত্যিকারের মাপ

এবার সত্যিকারের model. এই পর্বের জন্য নতুন একটি ছোট decoder বানানো হয়েছে: ১০ লাখের কিছু বেশি parameter, বাজারঘরের বানানো ডেলিভারি লাইনে character স্তরে ৬০০ ধাপ শেখানো। Attention-এর পুরো কৌশল কয়েকটি লাইনে (সংক্ষেপিত, আসল ফাইল থেকে নেওয়া):

q, k, v = self.qkv(x).chunk(3, dim=-1)     # x: শুধু নতুন token
q, k, v = [z.view(B, t, h, dk).transpose(1, 2) for z in (q, k, v)]
if past is not None:                       # এটুকুই আসল কৌশল
    k = torch.cat([past[0], k], dim=2)
    v = torch.cat([past[1], v], dim=2)
y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
return self.out(y.transpose(1, 2).reshape(B, t, -1)), (k, v)

দুটি জিনিস ভুল হলে নীরবে ফল নষ্ট হয়। এক, mask-এর আকার (t, off + t), যেখানে off হলো cache-এ থাকা token. দুই, position শুরু হবে off থেকে, শূন্য থেকে নয়। সত্যিকারের run (ডেস্কটপে; লেখার সময়ের সংস্করণ: Python 3.14.4, torch 2.15.0, numpy 2.5.4; সময় প্রতি token-এর, ২০টি token-এর median):

~/bazarghar-dl$ uv run python kv_generate.py
torch 2.15.0
step    1 loss 84.075
…
step  600 loss 0.327
params 1060480 | vocab 55
cuda greedy  identical tokens: True
cuda sampled identical tokens: True
text: অর্ডার 567: বরিশাল, চাল, দুই দিনে পৌঁছেছে। / অর্ডার 878: চট্টগ্রাম, সা
max |logit difference| cached vs full: 1.2e-05
cuda no cache ms/token  t=50: 4.99  t=200: 3.44  t=400: 3.74  t=800: 4.72  t=980: 10.04  | total 5.06 s
cuda cache    ms/token  t=50: 6.88  t=200: 5.78  t=400: 4.93  t=800: 5.29  t=980: 5.59  | total 5.60 s
cpu greedy  identical tokens: True
cpu sampled identical tokens: True
…
cpu no cache ms/token  t=50: 4.92  t=200: 14.52  t=400: 11.04  t=800: 23.42  t=980: 37.26  | total 19.16 s
cpu cache    ms/token  t=50: 1.43  t=200:  2.01  t=400:  2.06  t=800:  3.29  t=980:  2.48  | total  2.65 s
big params 26280448
big no cache ms/token  t=50: 19.81  t=200: 22.07  t=400: 35.57  t=800: 74.23  t=1180: 120.20  | total 71.01 s
big cache    ms/token  t=50: 24.07  t=200: 19.17  t=400: 19.23  t=800: 21.77  t=1180:  21.68  | total 24.71 s
big identical tokens: True
big prompt   250: first pass  21.57 ms | next token 21.81 ms
big prompt   500: first pass  41.94 ms | next token 18.98 ms
big prompt  1000: first pass  93.07 ms | next token 19.31 ms
big prompt  2000: first pass 255.28 ms | next token 29.28 ms
…
K shape (1, 4, 1000, 32) torch.float32 | cache bytes 4096000 | formula 2*L*d*T*bytes = 4096000

পড়ুন, আর সৎভাবে পড়ুন:

• মিল: Greedy আর sampling, দুটিতেই token এক। Logit-এর সর্বোচ্চ পার্থক্য ১.২e-05: float32-এর গোলের ঘরের পার্থক্য, “হুবহু” মানে এটুকুই।
• CPU-তে পরিষ্কার লাভ: t ≈ ৯৮০-তে ৩৭.২৬ ms বনাম ২.৪৮ ms; মোট ১৯.১৬ s বনাম ২.৬৫ s. একই লাইনে t = ২০০ আর ৪০০-র ওঠানামা দেখুন: মেশিনে অন্য কাজও চলছিল।
• ছোট model-এ GPU-তে কোনো লাভ নেই: এখানে cache-সহ মোট সময় বরং সামান্য বেশি, ৫.৬০ s বনাম ৫.০৬ s. প্রতি ধাপের খরচে GPU-তে কাজ পাঠানোর খরচই বড় (পর্ব ১২), matmul নয়। এই GPU সেই সময় অন্য কাজেও ব্যস্ত ছিল, তাই সংখ্যাটি ঢিলেঢালাভাবে পড়ুন।
• বড় model-এ (২.৬ কোটি parameter, না-শেখানো, শুধু সময়ের জন্য) লাভ ফিরে আসে: t ≈ ১,১৮০-তে ১২০.২০ ms বনাম ২১.৬৮ ms. Cache-এর সারি প্রায় সমতল।
• প্রথম token-এর দেরি: ২,০০০ token-এর prompt-এ প্রথম pass ২৫৫ ms, পরের token ২৯ ms. Cache প্রথম pass-কে ছোঁয় না।

এগুলো একটি local run-এর সংখ্যা, সাধারণ নিয়ম নয়। নিজের model-এ মাপুন।

৪ · Cache-এর memory — হাতে হিসাব

গুণ বাঁচে, দাম দেয় memory. প্রতি স্তরে K আর V, দুটি tensor, প্রতিটি (B, h, T, d/h). গুণ করলে h কাটা পড়ে:

KV memory = 2 · L · B · T · d · bytes

এখানে L স্তর, B batch, T token, d পুরো মাত্রা, আর bytes এক সংখ্যার মাপ (float32-এ ৪, float16-এ ২)। Head বাড়ালে cache বড় হয় না। পরীক্ষায় এই হিসাব প্রায়ই আসে, তাই একবার ধাপে ধাপে দেখুন।

এবার নানা মাপের model-এ। সংখ্যাগুলো কাল্পনিক উদাহরণ, কোনো নির্দিষ্ট পণ্যের নয়। প্রথম সারিটি আমাদের lab-এর model, যার সত্যিকারের cache আগের run-এ ৪০,৯৬,০০০ byte ছিল।

import numpy as np

SIZE = {"float32": 4, "float16": 2, "bf16": 2, "int8": 1}

def kv_bytes(L, h, d, T, dtype, B=1, kv_heads=None):
    """K আর V: 2 টি tensor, প্রতিটি (B, kv_heads, T, d/h), L স্তরে।"""
    return 2 * L * B * (kv_heads or h) * (d // h) * T * SIZE[dtype]

def human(n):
    for unit, s in ((2**30, "GiB"), (2**20, "MiB"), (2**10, "KiB")):
        if n >= unit: return f"{n / unit:7.1f} {s}"
    return f"{n:7d} B"

print(" L  h    d       T  dtype    B kv_h     memory")
for L, h, d, T, dt, B, kh in [(4, 4, 128, 1000, "float32", 1, None),
                              (6, 8, 512, 1024, "float16", 1, None),
                              (6, 2, 512, 1024, "float16", 1, None),
                              (6, 8, 512, 1024, "float16", 8, None),
                              (12, 12, 768, 1024, "float16", 1, None),
                              (32, 32, 4096, 8192, "float16", 1, None),
                              (32, 32, 4096, 128000, "float16", 1, None),
                              (32, 32, 4096, 128000, "float16", 1, 8)]:
    print(f"{L:2d} {h:2d} {d:4d} {T:7d} {dt:7s} {B:2d} {kh or h:4d} {human(kv_bytes(L, h, d, T, dt, B, kh))}")

L, h, d, T = 4, 4, 128, 1000
real = 0
for _ in range(L):                          # সত্যিকারের array বানিয়ে যাচাই
    K = np.zeros((1, h, T, d // h), np.float32); V = np.zeros_like(K)
    real += K.nbytes + V.nbytes
print("real arrays:", real, "| formula:", kv_bytes(L, h, d, T, "float32"))

 L  h    d       T  dtype    B kv_h     memory
 4  4  128    1000 float32  1    4     3.9 MiB
 6  8  512    1024 float16  1    8    12.0 MiB
 6  2  512    1024 float16  1    2    12.0 MiB
 6  8  512    1024 float16  8    8    96.0 MiB
12 12  768    1024 float16  1   12    36.0 MiB
32 32 4096    8192 float16  1   32     4.0 GiB
32 32 4096  128000 float16  1   32    62.5 GiB
32 32 4096  128000 float16  1    8    15.6 GiB
real arrays: 4096000 | formula: 4096000

Memory ফুরালে কী হয়? Cache নিজে কমে না, batch ছোট করতে হয়, বা অনুরোধ অপেক্ষায় থাকে। তাই serving system-গুলো একে সবচেয়ে দামি সম্পদের মতো সামলায়। এই প্রসঙ্গ আবার আসবে পর্ব ৫৯ (ছোট আর দ্রুত model) আর পর্ব ৮৩-এ (LLM serving).

৫ · Batched generation — কেন left padding

GPU-র লাভ পেতে কয়েকটি অনুরোধ একসঙ্গে চালাতে হয়। কিন্তু prompt-এর দৈর্ঘ্য আলাদা, আর tensor আয়তাকার। তাই পর্ব ১১-এর pad token লাগে। Classification-এ pad ডানে বসলেই চলে। Generation-এ নয়। (পর্ব ৩১-এর causal convolution-এর left padding অন্য জিনিস: সেখানে সময়ের অক্ষে শূন্য বসানো হয়, এখানে prompt সারিবদ্ধ করা।)

একটি ফাঁদ আছে: শুধু pad-এর দিক বদলালে হবে না। Pad-এর জন্য mask আর position id দুটোই লাগে। প্রথম তিনটি prompt, তারপর ৯৬টি prompt-এ, ১৬ করে batch-এ, সত্যিকারের run; প্রতি লাইনে দেখা হচ্ছে, batched ফল একা চালানোর সঙ্গে পুরো মেলে কি না:

~/bazarghar-dl$ uv run python padding_gen.py
prompt lengths: [18, 8, 15]
alone      prompt 2: '14: ঢাকা, চাল, দাম ঠিক ছ'
alone      prompt 3: 'েট, চাল, দাম ঠিক ছিল।\nঅর'
right ids  prompt 2: [18, 41, 54, 30, 47, 41, 2, 11, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
left  ids  prompt 2: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 18, 41, 54, 30, 47, 41, 2, 11]
right, logits[:, -1]        [True, True, False]
   prompt 3: 'ালেট, চাল, দাম ঠিক ছিল।\n'
left, positions 0..t-1      [True, False, False]
   prompt 3: 'েট, আলু, দাম ঠিক ছিল।\nঅর'
left, positions from mask   [True, True, True]
   prompt 3: 'েট, চাল, দাম ঠিক ছিল।\nঅর'
right, per-row bookkeeping  [True, True, True]
   prompt 3: 'েট, চাল, দাম ঠিক ছিল।\nঅর'
left, positions + KV cache  [True, True, True]
   prompt 3: 'েট, চাল, দাম ঠিক ছিল।\nঅর'
96 unseen-seed prompts, batches of 16: whole sequence equal to running alone
  right, logits[:, -1]         16 / 96
  left, positions 0..t-1       29 / 96
  left, positions from mask    96 / 96
  right, per-row bookkeeping   96 / 96
  left, positions + KV cache   96 / 96

শিক্ষাগুলো: (১) Right padding মাঝে মাঝে মিলে যায়, যেমন prompt ২-এ, আর ১৬টি sequence-এ। তিনটি উদাহরণ দেখে “ঠিক আছে” বলা বিপজ্জনক। সারি ৩-এ “সিল”-এর পরে এসেছে “ালেট”: অর্থহীন। (২) Left padding-এ position ভুল হলে “চাল”-এর জায়গায় “আলু”: পড়তে ভালো, কিন্তু ভুল, আর ৯৬-এর ৬৭টিতে এমন অমিল। Learned position embedding-এ এটি সরাসরি ধরা পড়ে। (৩) ঠিক পদ্ধতিতে ৯৬-এ ৯৬, cache-সহও। (৪) Per-row হিসাবও ৯৬-এ ৯৬: right padding নিজে অসম্ভব নয়; left padding সহজতর। Hugging Face transformers-এর decoder-only model-এ ডানের pad ধরলে সতর্কবার্তা আসে, আর সংশ্লিষ্ট আলোচনায় maintainer-রা left padding-এর কথাই বলেন (লেখার সময়ে, নভেম্বর ২০২৬)। শেখানোর সময় pad-এর দিক সাধারণত এত গুরুত্বপূর্ণ নয়, কারণ সেখানে প্রতিটি ঘরের loss আলাদা, আর mask ও ignore_index pad-কে বাদ দেয়।

৬ · FlashAttention-এর ধারণা — ফল একই, memory-র পথ আলাদা

KV cache কমায় হিসাবের সংখ্যা। FlashAttention (Dao ও সহকর্মী, ২০২২) অন্য জায়গায় হাত দেয়। এর প্রবন্ধের যুক্তি: attention ধীর শুধু গুণের কারণে নয়, GPU-র বড় memory আর চিপের ভেতরের ছোট দ্রুত memory-র মধ্যে ডেটা আনা-নেওয়ার কারণেও। Memory hierarchy-র ধারণা আগেই আছে কম্পিউটার আর্কিটেকচার কোর্সের পর্ব ৪৪-এ। সাধারণ attention T × T score-এর পুরো matrix বড় memory-তে লেখে, softmax করে, আবার পড়ে। অথচ শেষে লাগে শুধু output.

দুটি উপাদান। Tile: Q, K, V-কে ছোট টুকরোয় ভেঙে এক এক টুকরো দ্রুত memory-তে আনা। Online softmax: সব score একসঙ্গে না রেখেই softmax. এক query-র জন্য তিনটি সংখ্যা রাখলেই চলে: এখন পর্যন্ত সর্বোচ্চ score m, স্কেল করা যোগফল l, আর স্কেল করা output o. নতুন টুকরোর score s এলে:

m′ = max(m, max s) · l′ = l·e^m−m′ + ∑e^s−m′ · o′ = o·e^m−m′ + ∑e^s−m′·v

শেষে output = o / l. পর্ব ৩-এর stable softmax-এর মতো এখানেও সর্বোচ্চ মান বিয়োগ করে সংখ্যা নিরাপদ রাখা হয়। NumPy-তে যাচাই করুন, ফল কি সত্যিই একই?

import numpy as np

rng = np.random.default_rng(37)
T, d, b = 256, 32, 32                   # token, মাত্রা, এক টুকরোর মাপ
Q, K, V = rng.normal(size=(3, T, d))

def full(q, K, V):                      # সাধারণ পথ: T-টি score একসঙ্গে রাখা
    s = K @ q / np.sqrt(d)
    a = np.exp(s - s.max()); a /= a.sum()
    return a @ V, len(s)

def online(q, K, V, b):                 # টুকরো টুকরো: score জমে না, শুধু m, l, o
    m, l, o = -np.inf, 0.0, np.zeros(d)
    for i in range(0, len(K), b):
        s = K[i:i + b] @ q / np.sqrt(d)
        m2 = max(m, s.max())
        p = np.exp(s - m2)              # নতুন সর্বোচ্চ দিয়ে স্কেল করা
        r = np.exp(m - m2)              # আগের যোগফল নতুন স্কেলে আনা
        l = l * r + p.sum()
        o = o * r + p @ V[i:i + b]
        m = m2
    return o / l, b

a = np.array([full(q, K, V)[0] for q in Q])
for size in (1, 8, 32, 128):
    c = np.array([online(q, K, V, size)[0] for q in Q])
    print(f"tile {size:3d}: max |full - online| = {np.abs(a - c).max():.1e}")
print("scores held at once, all queries: full", T * T, "| online", b)

big = K[:4] * 40                        # খুব বড় score: সরাসরি e^s ফেটে যায় কি?
q = Q[0] * 40
s = big @ q / np.sqrt(d)
print("largest score", round(s.max(), 1), "| exp(score) =", np.exp(s.max()))
print("online on the same data:", np.round(online(q, big, V[:4], 2)[0][:3], 4))
print("full   on the same data:", np.round(full(q, big, V[:4])[0][:3], 4))

tile   1: max |full - online| = 8.9e-16
tile   8: max |full - online| = 8.9e-16
tile  32: max |full - online| = 8.9e-16
tile 128: max |full - online| = 7.8e-16
scores held at once, all queries: full 65536 | online 32
largest score 1716.3 | exp(score) = inf
online on the same data: [-1.466   1.5189 -0.3299]
full   on the same data: [-1.466   1.5189 -0.3299]

PyTorch-এ এটি আলাদা করে ইনস্টল করতে হয় না। F.scaled_dot_product_attention-এর docs (২.১৫, লেখার সময়ে, নভেম্বর ২০২৬) তিনটি বাস্তবায়নের কথা বলে: FlashAttention-2, Memory-Efficient Attention আর C++-এ লেখা সাধারণ গণিত। সবগুলো ডিফল্টে চালু, আর function নিজে input দেখে একটি বাছে। হাতের লেখা সংস্করণের সঙ্গে তুলনা, B = 4, h = 8, d_k = 64, causal mask, ৭–১১ বারের median; সত্যিকারের run:

~/bazarghar-dl$ uv run python sdpa_bench.py
torch 2.15.0
cpu  float32 T= 1024 hand   94.89 ms
             T= 1024 sdpa   57.96 ms | max diff 8.3e-07
cuda float32 T= 1024 hand   11.86 ms | extra peak  280.1 MiB
             T= 1024 sdpa    6.15 ms | extra peak   20.0 MiB | max diff 2.9e-06
cuda float32 T= 2048 hand   40.77 ms | extra peak 1064.1 MiB
             T= 2048 sdpa   22.67 ms | extra peak   48.0 MiB | max diff 2.3e-06
cuda float16 T= 1024 hand    4.65 ms | extra peak  144.1 MiB
             T= 1024 sdpa    1.14 ms | extra peak   10.0 MiB | max diff 2.0e-03
cuda float16 T= 2048 hand   19.18 ms | extra peak  536.1 MiB
             T= 2048 sdpa    4.37 ms | extra peak   24.0 MiB | max diff 2.0e-03
--- forcing one backend, GPU, float16, T=1024 ---
MATH                 mask       12.82 ms | extra peak 326.0 MiB
MATH                 is_causal  12.60 ms | extra peak 328.0 MiB
EFFICIENT_ATTENTION  mask        1.32 ms | extra peak  10.0 MiB
EFFICIENT_ATTENTION  is_causal   0.49 ms | extra peak   8.0 MiB
FLASH_ATTENTION      mask       refused: RuntimeError: No available kernel. Aborting execution.
FLASH_ATTENTION      is_causal   0.53 ms | extra peak   8.1 MiB

পড়ুন: সবচেয়ে বড় পার্থক্য সময়ে নয়, memory-তে। হাতের সংস্করণে T = ২,০৪৮, float32-এ অতিরিক্ত ১,০৬৪ MiB; SDPA-তে ৪৮ MiB. কারণ হাতের কোড 4 × 8 × T² আকারের score matrix বানায়, T দ্বিগুণে memory চার গুণ। GPU-তে সময়ের লাভ এখানে প্রায় দুই থেকে চার গুণ, dtype অনুযায়ী; CPU-তে প্রায় ১.৬ গুণ। ফলের পার্থক্য float32-এ ১০^−৬-এর ঘরে, float16-এ ২e-03: এগুলো ওই dtype-এর গোলের সীমা। মূল প্রবন্ধের গতির সংখ্যা তাদের নিজস্ব শর্তের; এই সংখ্যাগুলো একটি ব্যস্ত GPU-র একটি local run.

শেষ দুই ব্লক লক্ষ করুন। বুলিয়ান attn_mask-সহ FlashAttention জোর করলে এই build-এ প্রত্যাখ্যাত, is_causal=True-তে চলে। ডিফল্টে function অন্য kernel-এ চলে যায়, তাই সাধারণত ভুলটি চোখে পড়ে না। প্রতি সারির আলাদা pad mask-ও সাধারণ বুলিয়ান mask. তাই শুধু causal অংশটুকু হলে is_causal=True লেখা ভালো; নির্ভর করার আগে নিজের সংস্করণে পরীক্ষা করুন। MATH-এর memory বেশি কেন? Docs বলে, float16 input-এও ওই ব্যাকএন্ড মাঝের হিসাব float32-এ রাখে। আরেকটি ফাঁদ: এই function নিজে dropout বন্ধ করে না; model.eval()-এর সময় dropout_p শূন্য আপনাকেই পাঠাতে হয়।

৭ · Sliding window আর sparse attention — ফল বদলে খরচ কমানো

এতক্ষণের কৌশলে ফল বদলায়নি। এবার ভিন্ন লেনদেন: সবাইকে না দেখে কয়েকজনকে দেখা। Sliding window attention-এ প্রতিটি token শুধু শেষ W-টি token দেখে। Score-এর খরচ T²·d থেকে প্রায় T·W·d, আর cache-এর memory সীমিত থাকে W-তে। দাম? একটি স্তরে W-এর আগের কিছু দেখা যায় না। স্তরের পর স্তরে দৃষ্টি বাড়ে, পর্ব ৩১-এর receptive field-এর মতো।

import numpy as np, matplotlib.pyplot as plt
np.set_printoptions(precision=2, suppress=True, linewidth=72)
rng = np.random.default_rng(37)
T, d, W = 12, 8, 4                      # ১২ token, মাত্রা ৮, জানালা ৪
X = rng.normal(size=(T, d))
Wq, Wk, Wv = rng.normal(0, .4, (3, d, d))

def attn(x, window=None):
    s = (x @ Wq) @ (x @ Wk).T / np.sqrt(d)
    n = len(x); keep = np.tril(np.ones((n, n), bool))
    if window: keep &= ~np.tril(np.ones((n, n), bool), -window)
    s = np.where(keep, s, -np.inf)
    a = np.exp(s - s.max(-1, keepdims=True)); a /= a.sum(-1, keepdims=True)
    return x + a @ (x @ Wv), a

def stack(x, layers, window=None):
    for _ in range(layers): x = attn(x, window)[0]
    return x

yf, af = attn(X); yw, aw = attn(X, W)
print("per-token |full - window|:", np.abs(yf - yw).max(1))
X2 = X.copy(); X2[0] += 5.0                # প্রথম token (অর্ডার নম্বর) বদলানো
for layers in (1, 2, 3, 4):
    c_f = np.abs(stack(X2, layers) - stack(X, layers)).max(1)
    c_w = np.abs(stack(X2, layers, W) - stack(X, layers, W)).max(1)
    print(f"layers {layers}: token 0 reaches t = 11? full {c_f[11] > 0} | "
          f"window {c_w[11] > 0} | window reach {np.flatnonzero(c_w).max()}")
print("score multiplications, T = 12, d = 8: full", T * T * d, "| window",
      int(np.tril(np.ones((T, T)), 0).sum() - np.tril(np.ones((T, T)), -W).sum()) * d)
fig, ax = plt.subplots(1, 2, figsize=(6, 2.8))
for p, m, name in zip(ax, (af, aw), ("full causal", f"window W = {W}")):
    p.imshow(m, vmin=0, vmax=1, cmap="Blues"); p.set_title(name)
    p.set_xlabel("key position"); p.set_ylabel("query position")
fig.subplots_adjust(wspace=0.45, bottom=0.2); plt.show()

per-token |full - window|: [0.   0.   0.   0.   0.35 0.09 1.03 0.91 2.03 1.17 0.6  0.91]
layers 1: token 0 reaches t = 11? full True | window False | window reach 3
layers 2: token 0 reaches t = 11? full True | window False | window reach 6
layers 3: token 0 reaches t = 11? full True | window False | window reach 9
layers 4: token 0 reaches t = 11? full True | window True | window reach 11
score multiplications, T = 12, d = 8: full 1152 | window 336

Linformer আর Performer প্রকাশের সময় বেশ আলোচিত ছিল, আর তাদের লক্ষ্য T² ভাঙা। তবে লেখার সময়ে আমাদের জানা মতে, বড় LLM-এর মূলধারা এদের পথে যায়নি; exact attention, দক্ষ kernel আর cache নিয়েই চলেছে। বিষয়টি এখনো গবেষণাধীন, তাই “কখনো কাজে লাগবে না” ধরে নেবেন না। Attention পুরো সরিয়ে দেওয়ার কৌশল (যেমন state-space model) এই পর্বের বাইরে।

৮ · Long context — সীমা আর কৌশল

লম্বা context window-এর সীমা তিন জায়গায়: সময় (প্রথম pass), memory (cache), আর position. Position-এর সীমা নির্ভর করে পদ্ধতির ওপর। আমাদের learned position table-এর দৈর্ঘ্য ২,০৪৮; তার বাইরে গেলে সত্যিকারের run-এ দুই রকম ভুল:

max_len of the position table: 2048
CPU-তে:  IndexError: index out of range in self
GPU-তে:  AcceleratorError: CUDA error: device-side assert triggered

GPU-র বার্তাটি কোন ভুল, তা বলে না; তাই debug করতে আগে CPU-তে চালিয়ে দেখা কাজের কৌশল। RoPE-এ (পর্ব ৩৪) সীমাটি নরম: rotation-এর সূত্র যেকোনো দৈর্ঘ্যে হিসাব করা যায়, আর score-এ অবস্থানের প্রভাব নির্ভর করে শুধু আপেক্ষিক দূরত্বে। তবু দুটি কথা মনে রাখুন। বাস্তবায়নে প্রায়ই আগে থেকে হিসাব করা নির্দিষ্ট দৈর্ঘ্যের table থাকে। আর সূত্র চলা মানেই model সেই দৈর্ঘ্যে ভালো কাজ করবে, এমন নয়; শেখানোর দৈর্ঘ্যের বাইরে মান পড়তে পারে, যা পর্ব ৩৪-এর বিষয়।

Cache ছাঁটা আর আবার চালানো এক নয়. একটি সহজ কৌশল হলো cache-এর পুরনো অংশ ফেলে দেওয়া। কিন্তু রেখে দেওয়া K, V তৈরি হয়েছিল পুরনো context দেখে। সত্যিকারের run, ৪০০ token-এর পর W = ৬৪, একই পরের token:

cropped cache vs window rerun: max |logit diff| 3.443
cropped cache vs full context: max |logit diff| 33.388
window rerun  vs full context: max |logit diff| 31.866
same top token: True

তিনটি ভিন্ন ফল, যদিও এই উদাহরণে শীর্ষ token একই। Xiao ও সহকর্মীদের ২০২৩ সালের প্রবন্ধ দেখায়, শুধু শেষের K, V রাখলে text cache-এর চেয়ে লম্বা হলে ফল ভেঙে পড়ে, আর শুরুর কয়েকটি token-এর K, V রাখলে অনেকটা ফেরে। এদের নাম তাঁরা দেন attention sink. কারণ হিসেবে তাঁদের ব্যাখ্যা: শুরুর token-গুলো অর্থ নির্বিশেষে বেশি attention পায়।

Speculative decoding, একটি নোট. ছোট একটি model কয়েকটি token আগাম লেখে; বড় model এক pass-এ সেগুলো যাচাই করে, যতটুকু মেলে রাখে (Leviathan ও সহকর্মী, ২০২২–২৩; Chen ও সহকর্মী, ২০২৩)। Rejection sampling-এর নিয়ম মানলে output-এর বণ্টন বড় model-এর সমান থাকে। এখানে KV cache পিছিয়ে নেওয়া লাগে: বাতিল token-এর K, V ছেঁটে ফেলতে হয়। Beam search-এও (পর্ব ৫০) cache সারির ক্রম বদলাতে হয়।

৯ · AI-এর যুগে — generation কোড যাচাই

AI assistant generate লুপ মিনিটে লিখে দেয়। এই লুপের সবচেয়ে খারাপ ভুলগুলো ক্র্যাশ করে না, শুধু অন্য কিছু লেখে।

AI-এর লেখা generate লুপ হাতে পেলে — যাচাইয়ের চেকলিস্ট

• Cache ছাড়া আর cache-সহ, একই seed-এ, token মিলিয়েছি; logit-এর পার্থক্যও দেখেছি
• অনেক ভিন্ন দৈর্ঘ্যের prompt-এ batched ফল একা চালানোর সঙ্গে মিলিয়েছি; তিনটি উদাহরণে থামিনি
• K, V-এর shape আর মোট byte ছেপে হাতের সূত্রের সঙ্গে মিলিয়েছি
• model.eval(), torch.inference_mode() আর dropout_p=0 আছে
• সময় মেপেছি GPU-র synchronize-সহ, নিজের মেশিনে, কয়েকবার; প্রথম token আর পরের token আলাদা করে
• Backend, প্রবন্ধ আর সংখ্যা নিজে docs ও মূল লেখায় মিলিয়েছি; ব্যক্তিগত ডেটা বা key পাঠাইনি