اجرای مدل‌های زبانی بزرگ روی گوشی دیگر یک آزمایش تحقیقاتی نیست، بلکه یک واقعیت تجاری برای عرضه محصول است. با این حال، به محض اینکه از یک دموی ساده به یک محصول واقعی حرکت می‌کنید، مسئله تأخیر (latency) خود را دوباره نشان می‌دهد. شما مدل را کوچک کرده‌اید، وزن‌ها را کوانتیزه کرده‌اید، اما مرحله prefill همچنان کند است. جریان توکن‌ها متوقف می‌شود. رابط کاربری (UI) فریز می‌شود. مقصر به ندرت در جای درست قرار می‌گیرد.

در دستگاه‌های اندرویدی، محاسبات تقریباً هیچ‌گاه گلوگاه مرحله prefill در LLM نیست؛ بلکه پهنای باند حافظه است. SoCهای پرچمدار مدرن با هسته‌های قدرتمند GPU و NPU عرضه می‌شوند که می‌توانند محاسبات ریاضی را بسیار سریع‌تر از آنچه زیرسیستم حافظه می‌تواند تغذیه کند، انجام دهند. وقتی یک پیاده‌سازی ساده (naive) از attention را پروفایل می‌کنید، واحدهای اجرا اشباع نشده‌اند؛ آن‌ها در حال انتظار هستند. انتظار برای DRAM.

چرا مدل کوانتیزه شده شما هنوز کند به نظر می‌رسد

کوانتیزاسیون به اولین قدم پیش‌فرض برای استنتاج روی دستگاه (on-device inference) تبدیل شده است. کوچک کردن وزن‌ها از FP16 به INT8، اندازه مدل را نصف کرده و فضای ذخیره‌سازی را کاهش می‌دهد. این کار کمک می‌کند، اما تأخیر لایه attention را برطرف نمی‌کند. دلیل آن ساده است: کوانتیزاسیون مقدار داده‌ای را که ذخیره می‌کنید کاهش می‌دهد، اما تعداد تراکنش‌های حافظه‌ای را که مکانیزم attention انجام می‌دهد، کم نمی‌کند.

یک لایه standard multi-head attention که به روش کتابخانه‌ای پیاده‌سازی شده باشد، برای هر لایه سه بار رفت و برگشت کامل به DRAM دارد. ماتریس‌های Query، Key و Value از حافظه اصلی خوانده می‌شوند، امتیازها (scores) محاسبه می‌شوند و نتایج میانی دوباره نوشته می‌شوند. محاسبات ریاضی ساده هستند، اما جابجایی داده‌ها طاقت‌فرسا است. در اندروید، جایی که بودجه‌های توان و حرارتی محدود هستند، این الگو باعث آشفتگی (thrashing) باس حافظه می‌شود. پردازنده در واقع برای عبور از یک پل واحد، سه بار هزینه عابر می‌پردازد.

اگر مدل‌های INT8 را عرضه کرده‌اید و متعجب هستید که چرا مرحله prefill همچنان با افزایش طول پرامپت، به صورت درجه دوم (quadratically) مقیاس‌پذیر است، پاسخ همین است. وزن‌ها کوچک‌تر شده‌اند، اما ترافیک فعال‌سازی (activation traffic) همچنان عظیم است.

گلوگاه حافظه است، نه ریاضیات

برای درک راه حل، به مدل roofline نگاه کنید. پردازنده‌های گرافیکی (GPU) و NPUهای موبایل در تراشه‌هایی مانند Snapdragon 8 Gen 3 و Dimensity 9300، توان عملیاتی محاسباتی تئوری دارند که بسیار فراتر از توان رابط‌های LPDDR5X آن‌هاست. در یک کرنل attention ساده، هر هد (head) softmax را روی حاصل‌ضرب Q و K محاسبه می‌کند و سپس در V ضرب می‌کند. هر ماتریس امتیاز میانی در حافظه جهانی (global memory) مادی می‌شود. این یعنی اوج خواندن‌های DRAM شما با توان دوم طول توالی، یا O(n²) مقیاس می‌یابد. برای یک پرامپت ۱۰۲۴ توکنی، ترافیک حافظه به قدری زیاد است که زمان اجرا را تحت کنترل خود می‌گیرد.

هسته‌ها کمتر از حد بهینه استفاده می‌شوند زیرا نمی‌توانند تأخیر را پنهان کنند. پردازنده‌های مدرن برای پر نگه داشتن خط لوله‌ها (pipelines) به کش‌ها متکی هستند. وقتی یک الگوریتم مدام دچار cache miss می‌شود و داده‌ها را از DRAM فراخوانی می‌کند، واحدهای اجرا بیکار می‌مانند. هیچ میزان از کوانتیزاسیون نمی‌تواند این عدم تطابق ساختاری بین ظرفیت محاسباتی و عرضه حافظه را برطرف کند.

چگونه با استفاده از Tiling پهنای باند را بازیابی کنیم

راه حل، یک استراتژی Tiling است که امتیازهای میانی را به جای ارسال به DRAM، در SRAM روی تراشه نگه می‌دارد. این همان بینشی است که Flash Attention را هدایت می‌کند و برای پشته محاسباتی اندروید تطبیق داده شده است. به جای مادی کردن یک ماتریس امتیاز کامل n × n در حافظه، محاسبات را به تایل‌های (tiles) کوچکی تقسیم می‌کنید که در داخل کش L1 جا شوند. شما آمارهای محلی softmax را محاسبه می‌کنید، مقادیر max جاری و مجموع‌های نرمال‌سازی را انباشته می‌کنید و تنها خروجی‌های وزنی نهایی را به حافظه می‌نویسید.

این کار پیچیدگی پهنای باند را تغییر می‌دهد. اوج خواندن‌های DRAM از O(n²) به O(n کاهش می‌یابد، زیرا دیگر نیازی نیست ماتریس‌های امتیاز کامل را از طریق حافظه اصلی جابجا کنید. کارهای سنگین داخل SRAM، درست در کنار واحدهای اجرا انجام می‌شود.

برای یک مثال ملموس، اندازه تایل ۶۴ و ابعاد هد ۱۲۸ را در نظر بگیرید. تایل امتیاز ۱۶ کیلوبایت فضا اشغال می‌کند. این میزان اشغال فضا (footprint) به راحتی در کش L1 تراشه‌های پرچمدار فعلی مانند Snapdragon 8 Gen 3 و Dimensity 9300 قرار می‌گیرد. محاسبات ریاضی محلی باقی می‌ماند و باس حافظه نفس راحت می‌کشد.

پیاده‌سازی این مورد در اندروید

طرح الگوریتمی ساده است، اگرچه رعایت جزئیات اهمیت زیادی دارد.

Query، Key خود را تقسیم کنید...