পেজড অ্যাটেনশন (Paged Attention)-এর ম্যাজিক!

আমরা প্রতিদিন চ্যাটজিপিটি (ChatGPT) বা জেমিনির (Gemini) মতো যেসব লার্জ ল্যাঙ্গুয়েজ মডেল ব্যবহার করি, তার বেশিরভাগই প্রোপাইটারি বা পেইড; অর্থাৎ এগুলো কোম্পানিগুলোর নিজস্ব সার্ভারে চলে এবং আমরা সাবস্ক্রিপশন কিনে তা ব্যবহার করি। কিন্তু এর বাইরেও এক দারুণ জগত আছে—ওপেন সোর্স (Open-source) মডেল! সবচেয়ে মজার ব্যাপার হলো, এই ওপেন সোর্স মডেলগুলো চাইলে আমরা আমাদের নিজেদের ডিভাইসেই রান করতে পারি। এর জন্য কোনো বিশাল ডেটা সেন্টারের প্রয়োজন নেই, শুধু দরকার আমাদের নিজেদের কম্পিউটার আর একটি জিপিইউ (GPU) বা গ্রাফিক্স কার্ড। নিজেদের ডিভাইসে মডেল রান করে নতুন কিছু জেনারেট করার বা উত্তর তৈরি করার এই প্রক্রিয়াটিকেই বলা হয় 'এলএলএম ইনফারেন্সিং' (LLM Inferencing)।

User Prompt
     │
     ▼
Tokenizer
     │
     ▼
Transformer Model
     │
     ▼
KV Cache (stores previous token representations)
     │
     ▼
Next Token Generated

এখন, আমাদের আশেপাশে যেসব সাধারণ কনজ্যুমার গ্রাফিক্স কার্ড পাওয়া যায়, সেগুলোর একটি নির্দিষ্ট ক্ষমতা থাকে যে তারা প্রতি সেকেন্ডে কতগুলো টোকেন (শব্দের অংশ) জেনারেট করতে পারবে। কিন্তু যখনই আমরা একটু বড় পরিসরে এই লার্জ ল্যাঙ্গুয়েজ মডেলগুলো সার্ভ (serve) করার চেষ্টা করব, তখনই একটি অদ্ভুত সমস্যার মুখোমুখি হতে হয়। আপনি দেখবেন, আপনার কম্পিউট (Compute) বা প্রসেসিং ক্ষমতা শেষ হওয়ার অনেক আগেই মেমোরি (Memory) বা ভি-র‍্যাম (VRAM) শেষ হয়ে যাচ্ছে!

এখানে কিন্তু ব্যবহারকারী হিসেবে আপনার বা আমার কোনো দোষ নেই! আমরা সচরাচর নিজস্ব পিসিতে ইনফারেন্সের জন্য Ollama বা LM Studio-এর মতো যেসব দারুণ টুল ব্যবহার করি, সেগুলো মূলত 'সিঙ্গেল-ইউজার' বা ব্যক্তিগত ব্যবহারের কথা মাথায় রেখে তৈরি। এগুলো সাধারণত llama.cpp ব্যাকএন্ড ব্যবহার করে। মেমোরি বাঁচানোর জন্য এরা মডেলের সাইজকে কম্প্রেস (Quantization) করে GGUF ফরম্যাটে লোড করে এবং দরকার হলে জিপিইউ-এর পাশাপাশি কম্পিউটারের সাধারণ র‍্যামও (RAM) ব্যবহার করে। কিন্তু যখনই আপনি কোনো প্রজেক্টে বা ল্যাবে একসাথে অনেক ব্যবহারকারীকে (Multi-user) সার্ভ করতে যাবেন, তখন এই সাধারণ টুলগুলো হাঁপিয়ে ওঠে। কারণ তারা ব্যাকএন্ডে যে মেমোরি ম্যানেজমেন্ট পদ্ধতি ব্যবহার করে, তা বড় পরিসরে বা স্কেলে কাজ করার জন্য খুবই অদক্ষ (inefficient)।

চলুন, একজন প্রবলেম সলভারের মতো এই তিনটি ধাপ আগে বুঝে নিই!

মডেল লোডিং (Model Weights)

ইনফারেন্স শুরু করার আগে মডেলটিকে তার সমস্ত জ্ঞান নিয়ে জিপিইউ মেমোরিতে (VRAM) বসতে হয়। এই জ্ঞানগুলো হলো মডেলের 'প্যারামিটার' বা 'ওয়েট' (Weights)। ধরুন, আপনি যখন কোনো বই লেখেন, তখন পুরো বইয়ের থিম বা প্লট আগে আপনার মস্তিষ্কে লোড থাকতে হয়। ঠিক তেমনি, মডেলকেও কাজ শুরুর আগে তার বিশাল ডেটাসেটের নির্যাস মেমোরিতে লোড করে নিতে হয়, যা মেমোরির একটি বড় অংশ দখল করে।

কেভি ক্যাশের (KV Cache) আসল রহস্য: ট্রান্সফরমারের ম্যাথ থেকে মেমোরির ক্ষুধার গল্প!

চলুন আজ ট্রান্সফরমার মডেলের একদম হার্টের ভেতরে ঢুকে যাই। কেভি ক্যাশ বা পেজড অ্যাটেনশনের মতো ইঞ্জিনিয়ারিং জাদুগিরিগুলো কেন দরকার হলো, সেটা বুঝতে হলে আমাদের একটু পেছনের দিকে, অর্থাৎ ট্রান্সফরমার কীভাবে কাজ করে তার অংকেই ফিরে যেতে হবে। একজন গবেষকের চোখ দিয়ে দেখলে পুরো ব্যাপারটাই আসলে বিশাল বিশাল কিছু মেট্রিক্স আর ভেক্টরের গুণফল ছাড়া আর কিছুই নয়!

ট্রান্সফরমারের হার্টবিট: অ্যাটেনশন মেকানিজম (Attention Mechanism)

Input Tokens
     │
     ▼
Linear Projection
     │
 ┌───┼───────────────┐
 ▼   ▼               ▼
Query (Q)        Key (K)        Value (V)

        Q × Kᵀ
          │
          ▼
        Softmax
          │
          ▼
   Attention Weights
          │
          ▼
   Weighted Sum with V
          │
          ▼
Output Token Representation

ল্যাঙ্গুয়েজ মডেলের মূল ভিত্তি হলো "সেলফ-অ্যাটেনশন" (Self-Attention)। এর কাজ হলো বাক্যের প্রতিটি শব্দের সাথে অন্য সব শব্দের সম্পর্ক বা 'গুরুত্ব' বোঝা। এর সেই বিখ্যাত ইকুয়েশনটি হলো:

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

এখানে তিনটি ম্যাজিক্যাল ভেক্টর আছে: Q (Query), K (Key), এবং V (Value)। এই তিনটি জিনিস আসলে কী? চলুন একটা বাস্তব উদাহরণ দিয়ে বুঝি। ধরুন, আপনি আপনার কোনো নতুন রিসার্চ পেপারের জন্য লাইব্রেরিতে বই খুঁজছেন:

Query (Q): এটি হলো আপনার মাথার ভেতরের বর্তমান চাহিদা বা প্রশ্ন। যেমন, "আমার হিউম্যান-কম্পিউটার ইন্টারেকশনের (HCI) বই দরকার।" মডেল যখন নতুন কোনো টোকেন জেনারেট করতে যায়, সে একটি Query ভেক্টর তৈরি করে।

Key (K): এটি হলো বইয়ের শেলফে থাকা লেবেল বা ইনডেক্স ট্যাগ। যেমন, "কম্পিউটার সায়েন্স", "এআই", "ডিজাইন"। প্রতিটি আগের টোকেন মডেলকে বলে, "আমার কাছে এই এই প্রসঙ্গের ডেটা আছে।"

Value (V): এটি হলো বইয়ের ভেতরের আসল জ্ঞান বা কন্টেন্ট। মডেল যা করে তা হলো, বর্তমান শব্দের Query (Q)-এর সাথে পেছনের সব শব্দের Key (K)-এর গুণ বা ডট প্রোডাক্ট ($QK^T$) করে। যার সাথে যত বেশি মিল পাওয়া যায়, সেই শব্দের Value (V)-কে সে তত বেশি গুরুত্ব (Attention) দেয়।

কেভি ক্যাশ (KV Cache) এর জন্ম কেন হলো?

ধরে নিই মডেলটি "Artificial Intelligence is very..." এই বাক্যটির পর নতুন শব্দ জেনারেট করবে।

যদি সিস্টেমে কোনো ক্যাশ (Cache) না থাকে, তবে নতুন শব্দটি জেনারেট করার জন্য মডেলটিকে পেছনের ৪টি শব্দেরই Q, K, এবং V নতুন করে ক্যালকুলেট করতে হবে। এরপর যখন ৬ নম্বর শব্দটি জেনারেট করবে, তখন আবার পেছনের ৫টি শব্দের সব ভেক্টর শুরু থেকে হিসাব করতে হবে। ব্যাপারটা এমন দাঁড়াল যে, লাইব্রেরিতে প্রতিবার নতুন একটি বই খোঁজার জন্য আপনি লাইব্রেরির প্রতিটি শেলফের সমস্ত বইয়ের ট্যাগ এবং কন্টেন্ট একদম শুরু থেকে আবার পড়ছেন! এটি কম্পিউটেশনাল পাওয়ারের এক ভয়াবহ অপচয়।

সমাধান: বুদ্ধি করে আমরা আগের সব শব্দের Key (K) এবং Value (V) ভেক্টরগুলোকে মেমোরিতে সেভ করে রাখি। এটাই হলো KV Cache! যখন নতুন শব্দ জেনারেট হয়, মডেল শুধু বর্তমান শব্দটির জন্য একটি নতুন Query (Q) তৈরি করে। এরপর মেমোরি বা ক্যাশে আগে থেকেই সেভ করা থাকা পুরোনো K এবং V ভেক্টরগুলোর সাথে অংকটা কষে ফেলে। এতে প্রসেসরের বা কম্পিউটের বিশাল বাঁচা বেঁচে যায় এবং ইনফারেন্স স্পিড অনেক গুণ বেড়ে যায়।

Without KV Cache

Step 1:
Token1

Step 2:
Token1 + Token2
(recompute Token1)

Step 3:
Token1 + Token2 + Token3
(recompute Token1, Token2)

Step 4:
Token1 + Token2 + Token3 + Token4
(recompute everything again)

With KV Cache

Step 1:
Token1 → store KV

Step 2:
Token2 → store KV

Step 3:
Token3 → store KV

Step 4:
New Token
↓
Uses stored KV from memory
(no recomputation)

কিন্তু এটি এত মেমোরি খায় কেন?

Tokens generated →

Token1  Token2  Token3  Token4  Token5
  │       │       │       │       │
  ▼       ▼       ▼       ▼       ▼

KV Cache Memory Usage

█
██
███
████
█████

এখানেই আসে আসল ট্র্যাজেডি! আপনি হয়তো ভাবতে পারেন, কয়েকটা শব্দের Key আর Value সেভ করতে আর কত মেগাবাইটই বা লাগবে?

সমস্যা হলো, এগুলো আমাদের চেনা সাধারণ সংখ্যা নয়। এগুলো একেকটি বিশাল মাত্রার বা হাই-ডাইমেনশনাল ভেক্টর (High-dimensional vectors)।

হিসাবটা একটু প্র্যাকটিক্যালি ভাবি। ধরুন, আপনি Llama-2 এর মতো একটি মডেল রান করছেন:

• প্রতিটি টোকেন বা শব্দের জন্য অনেকগুলো নিউরাল নেটওয়ার্ক লেয়ার (যেমন ৪০টি লেয়ার) থাকে।
• প্রতিটি লেয়ারে আবার একাধিক 'অ্যাটেনশন হেড' (Attention Heads) থাকে।
• প্রতিটি হেডের জন্য K এবং V এর ভেক্টর ডাইমেনশন থাকে (যেমন ১২৮ ডাইমেনশন)।
• প্রতিটি সংখ্যা সেভ করতে মেমোরিতে ফ্লোটিং পয়েন্ট প্রিসিশন অনুযায়ী ২ বাইট (16-bit) জায়গা লাগে।

সব মিলিয়ে দেখা যায়, একটি মাঝারি সাইজের মডেলে মাত্র ১টি টোকেনের K এবং V সেভ করতে প্রায় ১ মেগাবাইট (MB) এর মতো মেমোরি বা VRAM দরকার হয়।

এখন ধরুন, আপনি একসাথে ১০০ জন ইউজারকে সার্ভ করছেন, এবং প্রত্যেকে গড়ে ২০০০ টোকেন জেনারেট করছে:

১০০ (ইউজার) × ২০০০ (টোকেন) × ১ MB = প্রায় ২০০ গিগাবাইট (GB)!

ভাবা যায়? শুধুমাত্র ইনফারেন্সের শর্ট-টার্ম মেমোরি বা ক্যাশ ধরে রাখতেই ২০০ জিবি মেমোরি উধাও! এর সাথে তো মডেলের নিজস্ব ওয়েটস (Weights) লোড করার জন্য মেমোরি আছেই।

আর ঠিক এই কারণেই, আগের আলোচনায় আমরা যেমনটা বলছিলাম—সনাতন পদ্ধতিতে এই বিশাল মেমোরি যদি আগে থেকেই বিশাল বড় ব্লক হিসেবে বরাদ্দ করে রাখা হয়, তবে তার বেশিরভাগই অব্যবহৃত পড়ে থেকে মেমোরির ভয়াবহ অপচয় (Fragmentation) ঘটায়। এই মেমোরি ম্যানেজমেন্টের গলার কাঁটা দূর করতেই 'পেজড অ্যাটেনশন' অপারেটিং সিস্টেমের পেজিং মেকানিজম ধার করে পুরো সিস্টেমকে বাঁচিয়ে দিয়েছে।

সহজ কথায়, প্রতিবার নতুন টোকেন জেনারেট করার সময় মডেলটি যেন পেছনের সব টোকেন নিয়ে নতুন করে অ্যাটেনশন (attention) হিসাব না কষে, সেজন্য সে ক্যাশ থেকে আগের ভ্যালুগুলো নিয়ে নেয়। এটি মডেলের এফিশিয়েন্সির জন্য খুব জরুরি। কেভি ক্যাশ না থাকলে, ১০০০ টোকেনের একটি রেসপন্স তৈরি করতে প্রতিটি ধাপে পেছনের সব টোকেনের হিসাব নতুন করে করতে হতো। তখন কাজের পরিমাণ $O(n)$ এর বদলে $O(n^2)$ হয়ে যেত! স্পিডের এই বিশাল পার্থক্যটা আপনি কাজ করলেই বুঝতে পারবেন।

মেমোরি অপচয়ের মূল সমস্যা

GPU Memory Layout

| User A Cache |      |
| User B Cache |      |
| User C Cache |      |
| User D Cache |      |

Empty spaces appear between allocations.
These gaps cannot be efficiently reused.

কিন্তু এখানে একটি বড় সমস্যা আছে। সনাতন পদ্ধতির কেভি ক্যাশ ইমপ্লিমেন্টেশনে প্রতিটি ইউজারের রিকোয়েস্টের জন্য আগে থেকেই মেমোরিতে একটি বড়, টানা বা (contiguous) ব্লক বরাদ্দ বা অ্যালোকেট (allocate) করে রাখা হয়।

কেন এটি মেমোরির অপচয়? একটু চিন্তা করুন তো!

ধরুন, রাজশাহীতে আপনি একটি রেস্টুরেন্টে খেতে গেছেন আপনার স্ত্রীকে নিয়ে। রেস্টুরেন্টের নিয়ম হলো, যে-ই আসুক না কেন, তাকে ১০ জনের একটি পুরো টেবিল বুক করতে হবে। আপনারা দুজন বসলেন, বাকি ৮টি চেয়ার ফাঁকা পড়ে রইল, অথচ বাইরে অন্য কাস্টমাররা অপেক্ষা করছেন। মেমোরির ক্ষেত্রেও ঠিক এই ঘটনাটাই ঘটে। আপনি একসাথে ১০০ জন ইউজারকে সার্ভ করছেন, কিন্তু আগে থেকে জানেন না কে কত বড় রেসপন্স চাইবে। তাই আপনি সর্বোচ্চ সাইজ ধরে মেমোরি বরাদ্দ করলেন—ধরলাম প্রতি রিকোয়েস্টের জন্য ২০৪৮ টোকেনের জায়গা। কিন্তু দেখা গেল গড়ে একেকজন ইউজার মাত্র ২০০ টোকেনের রেসপন্স পেয়েছে। আপনি কিন্তু আপনার আসল প্রয়োজনের চেয়ে ১০ গুণ বেশি মেমোরি আটকে রেখেছেন!

সমস্যা এখানেই শেষ নয়। প্রতিটি রিকোয়েস্ট মেমোরির একদম আলাদা বা আইসোলেটেড (isolated) ব্লক পায়। তাই ১০০ জন ইউজারের মধ্যে যদি ৮০ জন একই সিস্টেম প্রম্পট (System prompt) ব্যবহার করে, সনাতন পদ্ধতিতে ওই ৮০ জনের জন্যই আলাদা আলাদাভাবে একই প্রম্পটের ডেটা ডুপ্লিকেট হিসেবে সেভ হবে। এই মেমোরি শেয়ার করার কোনো উপায় নেই।

ফলে দেখা যায়, জিপিইউ মেমোরির মাত্র ২০-৩০% আপনি ঠিকমতো ব্যবহার করতে পারছেন। বাকিটা রিজার্ভ করা থাকলেও অব্যবহৃত পড়ে আছে। আর মেমোরি ছড়িয়ে ছিটিয়ে বা ফ্র্যাগমেন্টেড (fragmented) থাকায় মাঝখানের ফাঁকা জায়গাগুলোতে নতুন রিকোয়েস্ট ঢোকানোও যায় না। সিস্টেমের থ্রুপুট (throughput) কমিয়ে দেওয়ার মূল কারণ এটাই।

পেজড অ্যাটেনশন

Logical Token Blocks

Block1  Block2  Block3  Block4
   │       │       │       │
   ▼       ▼       ▼       ▼

Block Table (mapping)

   │       │       │       │
   ▼       ▼       ▼       ▼

Physical GPU Memory

| Block3 | Block1 | Block4 | Block2 |

এই বিশাল সমস্যার চমৎকার সমাধান হলো পেজড অ্যাটেনশন। আমরা যারা টেকনোলজি নিয়ে পড়াশোনা করেছি তারা অপারেটিং সিস্টেমের কোর্সে (OS) 'ভার্চুয়াল পেজিং মেমোরি' (virtual paging memory) সম্পর্কে পড়েছি। এখানে এই ধারণাটি ধার করেছে।

আপনি নিশ্চয়ই জানেন যে প্রোগ্রামগুলোকে র‍্যামের (RAM) বড় কোনো টানা অংশ দেওয়া হয় না। বরং OS মেমোরিকে ছোট ছোট নির্দিষ্ট সাইজের 'পেজ'-এ ভাগ করে ফেলে এবং ফিজিক্যাল মেমোরির যেকোনো জায়গায় ছড়িয়ে দেয়। তারপর একটি 'পেজ টেবিল' (page table) ব্যবহার করে এই ভার্চুয়াল অ্যাড্রেসের সাথে ফিজিক্যাল মেমোরির ম্যাপ তৈরি করে।

পেজড অ্যাটেনশন ঠিক এই কাজটিই করে কেভি ক্যাশের জন্য! চলুন দেখি কীভাবে:

ব্লক-লেভেল অ্যালোকেশন (Block-level allocation)

একটি রিকোয়েস্টের জন্য বিশাল বড় ব্লক বরাদ্দ না করে, কেভি ক্যাশকে ছোট ছোট নির্দিষ্ট সাইজের ব্লকে (সাধারণত ১৬ টোকেন) ভাগ করা হয়। এই ব্লকগুলো জিপিইউ মেমোরির যেকোনো জায়গায় থাকতে পারে, এদেরকে পাশাপাশি থাকতে হবে এমন কোনো কথা নেই।

বাস্তব উদাহরণ: এটি অনেকটা মাইনক্রাফট (Minecraft) গেমের ব্লকগুলোর মতো। গেমের ভেতরে একটি বিশাল ইমারত, যেমন লালবাগ কেল্লা বানাতে গেলে আপনাকে পুরো জায়গার জন্য একটি অখণ্ড বিশাল ব্লক ব্যবহার করতে হয় না। আপনি ছোট ছোট ব্লক দিয়ে যেকোনো জায়গায়, যেকোনো শেপে আপনার স্ট্রাকচার দাঁড় করাতে পারেন।

ব্লক টেবিল বা পেজ টেবিল (Block table)

প্রতিটি রিকোয়েস্টের জন্য একটি ব্লক টেবিল থাকে। এটি মূলত লজিক্যাল ব্লকের ইনডেক্সের সাথে জিপিইউ-এর ফিজিক্যাল ব্লকের একটি ম্যাপ। যখন মডেলের ৩৩ নম্বর টোকেনের ক্যাশ দরকার হয়, সে ব্লক টেবিলে খুঁজে দেখে ফিজিক্যাল মেমোরির কোথায় সেটি আছে এবং সেখান থেকে ফেচ (fetch) করে আনে। ফিজিক্যালি ব্লকটি কোথায় আছে, তা নিয়ে LLM-এর কোনো মাথা ব্যথা নেই।

শেয়ার্ড প্রিফিক্স (Shared prefixes)

User A Prompt
User B Prompt
User C Prompt

        │
        ▼

Shared Prefix KV Cache

        │
        ▼

Responses diverge

User A → answer A
User B → answer B
User C → answer C

আসল মজার খেলাটা এখানেই! যখন একাধিক ইউজার একই সিস্টেম প্রম্পট ব্যবহার করে, তখন সবার জন্য আলাদা করে ক্যাশ সেভ করার দরকার হয় না। সবার ব্লক টেবিল ফিজিক্যাল মেমোরির একই ব্লককে পয়েন্ট করে রাখে। যখন তাদের রেসপন্স আলাদা হতে শুরু করে (diverge), কেবল তখনই তাদের জন্য নতুন ব্লকের জায়গা বরাদ্দ করা হয়।

বাস্তব উদাহরণ: এটা অনেকটা একই বিশ্ববিদ্যালয়ে পড়া ছাত্রদের বাসের রুট শেয়ার করার মতো। একই ডিপার্টমেন্টের অনেক শিক্ষার্থী হয়তো একই বাসে ক্যাম্পাসে পৌঁছায় (শেয়ার্ড প্রম্পট), এরপর ক্লাস শেষে যার যার নির্দিষ্ট গন্তব্যে বা রুমে চলে যায় (ডাইভার্জ হওয়া)। প্রোডাকশনে যেখানে প্রায় সব রিকোয়েস্টেই কমন সিস্টেম প্রম্পট থাকে, সেখানে এই শেয়ারিংয়ের ফলে মেমোরি সেভিংস হয় অসাধারণ মাত্রায়!

এর প্রভাব কতটুকু?

পেজড অ্যাটেনশনের মূল পেপারে গবেষকরা মেপে দেখিয়েছেন যে, আগের সিস্টেমগুলো অ্যালোকেট করা মেমোরির মাত্র ২০-৩৮% ঠিকমতো ব্যবহার করতে পারতো। বাকিটা ফ্র্যাগমেন্টেশন আর ওভার-রিজার্ভেশনের কারণে নষ্ট হতো।

কিন্তু পেজড অ্যাটেনশন ব্যবহার করলে আপনি যা পাবেন:

• একই ল্যাটেন্সিতে (latency) আগের চেয়ে ২-৪ গুণ বেশি থ্রুপুট (Throughput)।
• মেমোরির অপচয় প্রায় শূন্যের কোঠায় নেমে আসে, কারণ যতটুকু ব্যবহার হয় ঠিক ততটুকুই বরাদ্দ করা হয়।
• ব্যাচিং (Batching) ভালো হওয়ার কারণে একই জিপিইউ হার্ডওয়্যারে অনেক বেশি রিকোয়েস্ট একসাথে সার্ভ করা যায়।

এ কারণেই vLLM-এর মতো ইনফারেন্স ইঞ্জিন, যা পেজড অ্যাটেনশনকে ভিত্তি করে বানানো, আজ প্রোডাকশন ডেপ্লয়মেন্টের জন্য সবার প্রথম পছন্দ। TensorRT-LLM এবং SGLang-এর মতো ফ্রেমওয়ার্কগুলোও এখন ফাস্ট ইনফারেন্সের জন্য একই ধরনের পেজিং মেকানিজম ব্যবহার করছে। কস্ট-পার-রিকোয়েস্ট (cost-per-request) কমানোর জন্য স্কেলে কাজ করার সময় এই অপটিমাইজেশনটি অপরিহার্য।