اتنشن تنک درونی (NSA در DeepSeek): مکانیزم، فرمول و پیاده‌سازی

پردازش بافت‌های طولانی (Long-context) یکی از بزرگ‌ترین چالش‌ها در استقرار و پیش‌آموزش مدل‌های زبانی مدرن است. در طول توالی‌های 64k یا بیشتر، لایه‌ی اتنشن کامل به‌تنهایی بخش عمده‌ای از زمان تاخیر دیکود (Decode latency) و حافظه‌ی کش کلید-مقدار (KV cache) را اشغال می‌کند؛ موضوعی که هزینه‌ی زیرساخت محاسباتی را به‌شدت بالا می‌برد.

معماری NSA (کوتاه‌شده‌ی Native Sparse Attention) که توسط تیم DeepSeek معرفی شد و جایزه‌ی بهترین مقاله‌ی ACL 2025 را به دست آورد، راه‌حلی برای این معضل است. این روش برخلاف تکنیک‌های پیشین تنک‌سازی که فقط در زمان استنتاج فعال می‌شدند، به‌صورت درونی در فرآیند پیش‌آموزش فعال است و با طراحی کرنل‌های هماهنگ با سخت‌افزار GPU، در بافت‌های طولانی سرعتی بالاتر از FlashAttention فراهم می‌کند.

در این درس ساختار ریاضی و معماری سه شاخه‌ی موازی NSA (شاخه‌ی فشرده، انتخابی و پنجره‌ی لغزان) را بررسی می‌کنیم و نحوه‌ی پیاده‌سازی و ترکیب وزن‌دار آن‌ها با گیت یادگرفتنی را به‌شکل گام‌به‌گام یاد می‌گیریم.

در بافت 64k توکن، لایه‌ی اتنشن (attention) حدود 70 تا 80 درصد از تاخیر دیکود (decode latency) را می‌بلعد. تمام آزمایشگاه‌های مدل‌های متن‌باز برنامه‌ای برای رفع این مشکل دارند. در این میان، ایده‌ی NSA متعلق به DeepSeek (برنده‌ی جایزه‌ی بهترین مقاله‌ی ACL 2025) همان راهکاری است که ماندگار شد: سه branch موازی اتنشن — توکن‌های فشرده‌شده و درشت‌دانه (coarse-grained)، توکن‌های ریزدانه (fine-grained) که به‌صورت انتخابی نگه داشته شده‌اند، و پنجره‌های لغزان (sliding windows) برای بافت موضعی — که با یک گیت یادگرفتنی ترکیب می‌شوند. این روش با سخت‌افزار هماهنگ است (kernel-friendly)، به‌صورت درونی قابل آموزش است (در مرحله‌ی پیش‌آموزش عمل می‌کند، نه این‌که صرفاً هنگام استنتاج وصله‌پینه شده باشد) و در دیکود توالی‌های 64k، سریع‌تر از FlashAttention اجرا می‌شود و کیفیتی معادل یا بهتر از اتنشن کامل ارائه می‌دهد. در این درس، هر سه branch را از ابتدا تا انتها پیاده‌سازی می‌کنیم و نشان می‌دهیم چرا تنک‌سازی (sparsity) آن در کل مسیر به‌صورت انتها‌به‌انتها مشتق‌پذیر است.

اهداف یادگیری

• نام بردن سه branch اتنشن در NSA و توصیف نقشی که هرکدام ثبت می‌کنند.
• توضیح این‌که چرا NSA «به‌صورت درونی قابل آموزش» است، در حالی که روش‌های پیشین اتنشن تنک (sparse attention) فقط در زمان استنتاج به کار می‌رفتند.
• محاسبه‌ی میزان صرفه‌جویی در محاسبات اتنشن با NSA در مقایسه با اتنشن کامل در بافت 64k، بر حسب اندازه‌ی بلوک فشرده‌سازی و مقدار top-k انتخابی.
• پیاده‌سازی ترکیب سه branch با پایتون استاندارد روی یک توالی ساختگی کوتاه و راستی‌آزمایی رفتار وزن‌های گیت.

صورت مسئله

اتنشن کامل در طول توالی N به ازای هر لایه، زمان O(N^2) و کش کلید-مقدار (KV cache) به میزان O(N) مصرف می‌کند. در 64k توکن، ارقام مربوط به محاسبات و پهنای باند حافظه فاجعه‌بار می‌شوند. برآورد تئوریِ ارائه‌شده در مقاله‌ی NSA نشان می‌دهد که در 64k توکن، اتنشن حدود 70 تا 80 درصد از کل تاخیر دیکود را تشکیل می‌دهد. تمام متغیرهای پسین — مانند TTFT، نرخ توکن بر ثانیه و هزینه‌ی پردازش هر یک میلیون توکن — تحت‌الشعاع هزینه‌ی اتنشن قرار دارند.

اتنشن تنک راه‌حل طبیعی و بدیهی این مشکل است. تلاش‌های پیشین (prior attempts) به دو دسته تقسیم می‌شوند: تنک‌سازی با الگوی ثابت (پنجره‌ی لغزان، استراید یا بلوک‌های موضعی) که اطلاعات را دور می‌ریزد و در وظایف بازخوانی (recall) دوربرد شکست می‌خورد؛ دسته‌ی دوم، تنک‌سازی در زمان استنتاج است (هرس کردن کش کلید-مقدار، H2O، StreamingLLM) که روی مدلی اعمال می‌شود که با اتنشن متراکم پیش‌آموزش دیده است. این روش تنها بخش کوچکی از شتاب بالقوه را محقق می‌کند، زیرا در زمان آموزش هرگز از مدل خواسته نشده بود که اطلاعات را از میان الگوی تنک هدایت کند.

اتنشن تنک درونی یا NSA (توسط Yuan و همکاران از DeepSeek ،PKU و UW، برنده‌ی جایزه‌ی بهترین مقاله‌ی ACL 2025، با شناسه‌ی arXiv:2502.11089) هر دو مشکل را حل می‌کند: الگوی تنکی که مدل در طول پیش‌آموزش یاد می‌گیرد، در قالب الگوریتمی هماهنگ با کرنل پیاده می‌شود که صرفه‌جویی محاسباتی واقعی را هنگام استنتاج به همراه می‌آورد. تا دو سال آینده، NSA یا نوادگان مستقیم آن، اتنشن پیش‌فرض در تمام مدل‌های پیشروی بافت‌بلند خواهند بود.

مفاهیم کلیدی

سه branch موازی

برای هر کوئری (query)، روش NSA لایه‌ی اتنشن را سه بار، روی سه نمای متفاوت از کش کلید-مقدار اجرا می‌کند:

1. branch فشرده‌شده (Compressed branch): توکن‌ها در بلوک‌هایی با اندازه‌ی l (معمولاً 32 یا 64) گروه‌بندی می‌شوند. هر بلوک از طریق یک MLP یادگرفتنی کوچک، به یک توکن خلاصه‌ی واحد فشرده می‌شود. کوئری روی این توکن‌های فشرده‌شده اتنشن اعمال می‌کند و نمایی درشت‌دانه از کل توالی به دست می‌آورد.

2. branch انتخابی (Selected branch): با استفاده از امتیازهای اتنشنِ به دست آمده از branch فشرده، top-k بلوکی که بیشترین ارتباط را با کوئری فعلی دارند مشخص می‌شوند. سپس توکن‌های ریزدانه (فشرده‌نشده) از آن بلوک‌ها خوانده شده و کوئری روی همه‌ی آن‌ها اتنشن اعمال می‌کند. اتنشن branch فشرده را مثل سیگنال مسیریابی برای انتخاب بلوک‌ها در نظر بگیرید.

3. branch پنجره‌ی لغزان (Sliding-window branch): کوئری برای ثبت بافت موضعی، روی جدیدترین توکن‌ها به تعداد W (معمولاً 512) اتنشن اعمال می‌کند. این branch الگوهای کوتاه‌برد و به‌شدت ساختارمند (دستور زبان، مرجع‌گزینی موضعی) را که دو branch دیگر ممکن است از قلم بیندازند، ثبت می‌کند.

خروجی‌های این سه branch از طریق یک گیت یادگرفتنی به ازای هر موقعیت ترکیب می‌شوند:

$$\text{out} = g_{\text{cmp}} \cdot \text{out}_{\text{cmp}} + g_{\text{sel}} \cdot \text{out}_{\text{sel}} + g_{\text{win}} \cdot \text{out}_{\text{win}}$$

مقادیر g_cmp, g_sel, g_win وزن‌های گیت هستند که از یک MLP کوچک روی کوئری به دست می‌آیند. لزومی ندارد مجموع آن‌ها 1 شود؛ این وزن‌ها می‌توانند branchها را مستقل از هم وزن‌دهی کنند.

چرا این روش «به‌صورت درونی قابل آموزش» است؟

گام انتخاب بلوک‌ها (انتخاب top-k) گسسته است و عملیات گسسته جریان گرادیان (gradient) را قطع می‌کنند. کارهای قبلی در اتنشن تنک یا پس‌انتشار (backprop) را از مرحله‌ی انتخاب عبور نمی‌دادند (که آموزش را محدود می‌کرد)، یا از تقریب‌های پیوسته استفاده می‌کردند که در زمان استنتاج تنک‌سازی واقعی ایجاد نمی‌کرد.

روش NSA این مانع را دور می‌زند: اتنشن در branch فشرده، خود یک اتنشن درشت‌دانه‌ی مشتق‌پذیر روی کل توالی است. عملیات top-k صرفاً بالاترین امتیازهای اتنشن از branch فشرده را دوباره به کار می‌گیرد تا مشخص کند کدام بلوک‌های ریزدانه بارگذاری شوند. گرادیان‌ها از طریق خروجیِ خودِ branch فشرده به امتیازهای آن می‌رسند، نه از مسیر انتخاب بلوک‌ها. درون بلوک‌های انتخاب‌شده نیز، اتنشن روی کلیدها و مقادیر منتخب، یک اتنشن معمولی است؛ بنابراین سهم آن‌ها در خروجی نهایی هم مشتق‌پذیر خواهد بود. خودِ انتخاب اندیس‌های top_k گرادیانی دریافت نمی‌کند، بلکه فقط تعیین می‌کند کدام بلوک‌ها از حافظه بارگذاری شوند.

به همین دلیل است که NSA می‌تواند در کل مسیر پیش‌آموزش به‌شکل انتها‌به‌انتها استفاده شود. مدل یاد می‌گیرد اطلاعات را به‌طور هم‌زمان از میان این سه branch هدایت کند و الگوی تنکی می‌سازد که در زمان استنتاج، شتاب موعود را واقعاً محقق می‌کند.

کرنل هماهنگ با سخت‌افزار

کرنل NSA برای سلسله‌مراتب حافظه‌ی GPUهای مدرن طراحی شده است. این کرنل، کوئری‌ها را بر اساس گروه‌های GQA بارگذاری می‌کند (حلقه‌ی بیرونی)، بلوک‌های کلید-مقدار تنک مربوط به هر گروه را واکشی می‌کند (حلقه‌ی درونی) و اتنشن را روی SRAM به اجرا درمی‌آورد. از آنجا که تمام کوئری‌های یک گروه، بلوک‌های انتخابی یکسانی را می‌بینند (انتخاب بلوک به ازای هر گروه کوئری انجام می‌شود، نه به ازای هر سرِ کوئری)، هزینه‌ی بارگذاری کلید-مقدار بین اعضای گروه سرشکن می‌شود؛ در نتیجه شدت محاسباتی (arithmetic intensity) بالا می‌ماند.

این مقاله گزارش می‌دهد که کرنل‌های Triton آن در دیکود بافت‌های 64k تا 9 برابر سریع‌تر از FlashAttention اجرا می‌شوند و این ضریب شتاب با افزایش طول توالی بیشتر هم می‌شود. کرنل‌های هر دو مسیر رفت و برگشت (forward و backward) در دسترس قرار گرفته‌اند.

بودجه‌ی محاسباتی

فرض کنید N طول توالی، l اندازه‌ی بلوک فشرده‌سازی، k تعداد انتخاب top-k، و w پنجره‌ی لغزان و b اندازه‌ی بلوک انتخابی باشد (که معمولاً برابر با l است).

• branch فشرده: به ازای هر کوئری O(N/l) کلید، یعنی در مجموع O(N * N / l).
• branch انتخابی: به ازای هر کوئری O(k * b) کلید، یعنی O(N * k * b).
• branch لغزان: به ازای هر کوئری O(w) کلید، یعنی O(N * w).

مجموع: O(N * (N/l + k*b + w)).

با مقادیر N = 64k, l = 64, k = 16, b = 64, w = 512: هزینه‌ی هر کوئری برابر با 1000 + 1024 + 512 = 2536 keys کلید است، در حالی که اتنشن کامل به 64000 keys نیاز دارد؛ یعنی 25 برابر کاهش محاسبات.

با مقادیر N = 128k, l = 64, k = 16, b = 64, w = 512: هزینه‌ی هر کوئری برابر با 2000 + 1024 + 512 = 3536 keys می‌شود، در حالی که اتنشن کامل 128000 keys می‌خواهد. این یعنی 36 برابر کاهش محاسبات. این مزیت با بلندتر شدن طول توالی افزایش می‌یابد که هدف اصلی طراحی هم دقیقاً همین بوده است.

مقایسه‌ی روش‌ها

روش MoBA (از شرکت Moonshot، با شناسه‌ی arXiv:2502.13189) هم‌زمان منتشر شد و رویکرد مشابه «سه تا بهتر از یکی است» را پیش گرفت و ایده‌ی MoE را روی بلوک‌های اتنشن پیاده کرد. NSA و MoBA دو معماری اصلی هستند که باید برای پیش‌آموزش مدل‌های بافت‌بلند در سال 2026 بشناسید.

پیاده‌سازی عملی

فایل code/main.py هر سه branch را روی یک توالی ساختگی کوتاه پیاده کرده و بخش‌های زیر را نشان می‌دهد:

• بخش فشرده‌ساز MLP (برای سادگی آموزش از خط مبنای میانگین‌گیری استخری استفاده شده است؛ در NSA واقعی از یک MLP یادگرفتنی استفاده می‌شود).
• انتخاب بلوک‌های top-k که با امتیازهای branch فشرده هدایت می‌شود.
• اتنشن پنجره‌ی لغزان روی آخرین w توکن.
• ترکیب وزن‌دار با گیت.
• چاپ شمارش محاسبات در مقایسه با اتنشن کامل.

گام 1: فشرده‌سازی توکن‌ها در قالب بلوک‌ها

def compress(K, l):
    n = len(K)
    n_blocks = (n + l - 1) // l
    out = []
    for b in range(n_blocks):
        start, end = b * l, min((b + 1) * l, n)
        block = K[start:end]
        summary = [sum(row[d] for row in block) / len(block) for d in range(len(K[0]))]
        out.append(summary)
    return out

گام 2: اتنشن در branch فشرده

اتنشن سافت‌مکس (Softmax) کوئری را روی کلیدهای فشرده‌شده اجرا کنید. امتیازهای این branch فشرده هم‌زمان نقش سیگنال راهنما برای انتخاب top-k را ایفا می‌کنند.

گام 3: انتخاب بلوک‌های top-k

اندیس‌های k بلوک فشرده‌ای را که بالاترین امتیاز را آورده‌اند انتخاب کنید. توکن‌های اصلی و فشرده‌نشده‌ی آن بلوک‌ها را بارگذاری کرده و اتنشن را روی آن‌ها اجرا کنید.

گام 4: اتنشن پنجره‌ی لغزان

آخرین w توکن را بردارید و اتنشن استاندارد را روی آن‌ها اجرا کنید.

گام 5: گیت و ترکیب نهایی

یک MLP کوچک روی کوئری، سه وزن گیت تولید می‌کند. خروجی نهایی، مجموع وزن‌دارِ خروجی‌های این سه branch است.

گام 6: شمارش محاسبات

تعداد کلیدهایی را که به ازای هر کوئری در هر branch دیده شده‌اند، همراه با مجموع آن‌ها چاپ کنید. این رقم را با N (اتنشن کامل) مقایسه کنید. روی یک نمونه‌ی ساختگی 1024 توکنی با مقادیر l = 32, k = 4, w = 128، روش NSA به ازای هر کوئری 32 + 128 + 128 = 288 کلید را پردازش می‌کند، در حالی که اتنشن کامل به 1024 کلید نیاز دارد — یعنی 3.5 برابر کمتر.

کاربرد در محیط عملیاتی

روش NSA در پایپ‌لاین پیش‌آموزش بافت‌بلندِ خود DeepSeek به کار گرفته شده است. وضعیت یکپارچه‌سازی در ابزارهای استنتاج عمومی تا آوریل 2026:

• سیستم داخلی DeepSeek: بومی و پیش‌فرض؛ وزن‌های منتشرشده از NSA یا جانشین آن یعنی DSA (اتنشن تنک DeepSeek) استفاده می‌کنند.
• vLLM: پشتیبانی آزمایشی از NSA برای وزن‌های سری DeepSeek-V3.x در حال توسعه است.
• SGLang: بنچمارک‌های NSA منتشر شده است؛ مسیر استفاده‌ی نهایی از پایپ‌لاین vLLM پیروی می‌کند.
• llama.cpp / CPU: پشتیبانی نمی‌شود؛ بار پردازشی ناشی از تفکیک کرنل با نرخ پردازش (throughput) پردازنده‌ها توجیه اقتصادی ندارد.

چه زمانی باید سراغ NSA بروید:

• پیش‌آموزش یا ادامه‌ی آموزش مدلی با بافت 64k به بالا در شرایطی که بودجه‌ی محاسباتی جدی دارید.
• استنتاج چک‌پوینت‌های بافت‌بلند خود DeepSeek؛ چرا که وزن‌های آن بر پایه‌ی ساختار درونی NSA آموزش دیده‌اند.

چه زمانی نباید سراغ آن بروید:

• سرویس‌دهی به مدلی که پیش‌تر با اتنشن متراکم پیش‌آموزش دیده است؛ چرا که نمی‌توانید بدون ادامه‌ی آموزش، NSA را به آن اضافه کنید.
• بافت‌های زیر 16k؛ چون بار محاسباتی ناشی از مدیریت سه branch بیشتر از صرفه‌جویی آن خواهد بود.
• چت‌های تعاملی با batch size برابر 1؛ هرچند دیکودِ حساس به تاخیر از آن سود می‌برد، اما این مزیت فقط در بافت‌های طولانی ملموس است.

ارائه‌ی نهایی محصول

این درس فایل outputs/skill-nsa-integrator.md را تولید می‌کند. این ابزار با دریافت مشخصات یک پروژه پیش‌آموزش بافت‌بلند، یک برنامه‌ی یکپارچه‌سازی NSA ارائه می‌دهد: اندازه‌ی بلوک فشرده‌سازی، مقدار top-k، پنجره‌ی لغزان، عرض شبکه MLP گیت، انتخاب نوع کرنل، و ارزیابی‌های اختصاصی بافت‌بلند که تغییر معماری را توجیه می‌کنند.

تمرین‌ها

1. اسکریپت code/main.py را روی یک توالی ساختگی 1024 توکنی اجرا کنید. متغیرهای (l, k, w) را روی سه حالت پیش‌فرض مختلف تغییر دهید و تعداد محاسبات را چاپ کنید. حالتی را بیابید که کمترین تعداد کلید در هر کوئری را ثبت کند، در حالی که در تست «پیدا کردن سوزن در انبار کاه»، نرخ بازخوانی 95 درصدی را در مقایسه با اتنشن کامل حفظ کند.

2. ماژول فشرده‌ساز میانگین‌گیر را با یک MLP یادگرفتنیِ کوچک (2 لایه با بعد مخفی 32) جایگزین کنید. آن را روی داده‌های ساختگی که سیگنال اصلی در آن‌ها میانگین یک بلوک است آموزش دهید. سپس اختلاف پرپلکسیتی (perplexity) را روی داده‌های آزمون نسبت به خط مبنای میانگین‌گیر بسنجید.

3. شبکه‌ی MLP مربوط به گیت را پیاده‌سازی کنید. این شبکه کوئری را به عنوان ورودی می‌گیرد و سه اسکالر خروجی می‌دهد. نشان دهید که گیت منطقی رفتار می‌کند: در کوئری‌های تصادفی وزن‌ها تقریباً یکنواخت باشند و زمانی که کوئری به بلوکی در دوردست‌ها ارجاع می‌دهد، وزن بالایی به branch انتخابی بدهد.

4. بودجه‌ی حافظه‌ی کش کلید-مقدار را برای یک مدل 70B مجهز به NSA در بافت 128k محاسبه کنید. پارامترها: 8 سرِ کلید-مقدار، بعد هر سر 128 و فرمت BF16. خروجی را با اتنشن کامل و معماری MLA (که در فاز 10 · 14 بررسی شد) مقایسه کنید. طولی از توالی را پیدا کنید که در آن، کش کلید-مقدارِ branch ریزدانه‌ی NSA با اتنشن کامل برابر می‌شود.

5. بخش 4 مقاله‌ی NSA (شناسه‌ی arXiv:2502.11089) را بخوانید و در سه جمله توضیح دهید چرا امتیازهای اتنشنِ branch فشرده دوباره برای انتخاب top-k استفاده می‌شوند و امتیاز مسیریابیِ جداگانه‌ای محاسبه نمی‌شود. پاسخ خود را با سازوکار جریان گرادیان پیوند دهید.

اصطلاحات کلیدی

برای مطالعه‌ی بیشتر

• Yuan و همکاران — Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention (arXiv:2502.11089, ACL 2025 Best Paper) — مقاله‌ی اصلی
• گزارش فنی DeepSeek-V3 (arXiv:2412.19437) — خانواده‌ی معماری که NSA برای آن طراحی شد
• Moonshot AI — MoBA: Mixture of Block Attention for Long-Context LLMs (arXiv:2502.13189) — پژوهش هم‌زمان با تمرکز بر اتنشن بلوکی به سبک MoE
• Beltagy و همکاران — Longformer: The Long-Document Transformer (arXiv:2004.05150) — خاستگاه پنجره‌ی لغزان
• Xiao و همکاران — StreamingLLM: Efficient Streaming Language Models with Attention Sinks (arXiv:2309.17453) — خط مبنای تنک‌سازی در استنتاج با سینک اتنشن (attention sink) که NSA آن را ارتقا داد
• Dao و همکاران — FlashAttention-2 (arXiv:2307.08691) — خط مبنای اتنشن کامل که کرنل‌های NSA در بافت 64k آن را پشت سر می‌گذارند