PRISM: Breaking the O(n) Memory Wall in Long-Context LLM Inference via O(1) Photonic Block Selection

1. بيان المشكلة: جدار عرض نطاق الذاكرة (Memory Bandwidth Wall)

تتعرض عملية الاستدلال الحالية في النماذج اللغوية الكبيرة (LLMs) لعائق ليس ناتجاً عن القدرة الحسابية (FLOPs)، بل عن عرض نطاق الذاكرة.

• عنق زجاجة ذاكرة الـ KV Cache: أثناء عملية فك التشفير التكراري (Autoregressive Decoding)، تتطلب توليد كل توكن جديد قراءة كامل ذاكرة الـ KV cache المتراكمة من جميع التوكنات السابقة. بالنسبة للسياقات الطويلة (مثل 128 ألف توكن أو أكثر)، يمكن أن تتجاوز ذاكرة الـ KV cache حاجز الـ 40 جيجابايت.
• مشكلة التوسع من رتبة $O(n)$: الطرق الإلكترونية التقليدية (حتى تلك التي تستخدم أساليب اختيار على مستوى الكتل مثل Quest أو RocketKV) يجب أن تظل تقوم بـ مسح (Scan) لجميع $N$ من كتل الذاكرة لحساب درجات التشابه واختيار أفضل $k$ مرشحين. يؤدي هذا إلى تكلفة قراءة ذاكرة من رتبة $O(N)$ لكل خطوة فك تشفير، وهي تكلفة تتزايد خطياً مع طول السياق.
• السياق الصناعي: حتى بنية "Vera Rubin" القادمة من NVIDIA تقر بهذا الأمر عبر تخصيص وحدة معالجة بيانات (DPU) مخصصة لإدارة KV cache، مما يؤكد أن الوصول إلى الذاكرة، وليس العمليات الحسابية، هو القيد الأساسي.

2. المنهجية: بنية PRISM

يقترح المؤلفون PRISM (التصنيف الضوئي عبر التشابه الداخلي مع أوزان الحلقات الدقيقة - Photonic Ranking via Inner-product Similarity with Microring weights)، وهو محرك ضوئي متخصص مصمم للقضاء على مسح الذاي من رتبة $O(N)$ تماماً.

المفهوم الجوهري: البث والوزن الضوئي (Broadcast-and-Weight - B&W)

يستفيد PRISM من الفيزياء الفريدة للبصريات لأداء البحث عن التشابه في زمن انتقال ضوئي قدره $O(1)$، بغض النظر عن عدد المرشحين ($N$).

• التطابق الهيكلي: مهمة اختيار كتل الـ KV هي "بحث عن التشابه" حيث يتم مقارنة متجه استعلام واحد مع $N$ من التواقيع الثابتة. وهذا يتوافق تماماً مع نموذج البث والوزن (Broadcast-and-Weight):
  1. البث (Broadcast): يتم تشفير متجه الاستعلام على $d$ من الأطوال الموجية وتوزيعه سلبياً على جميع $N$ من القنوات المرشحة في آن واحد (لا يوجد وصول تسلسلي للذاكرة).
  2. الوزن (Weighting): تُخزن تواقيع كتل الـ $N$ كأوزان قابلة للبرمجة في الرنانات الحلقية الدقيقة (MRRs).
  3. التكامل (Integration): تقوم الكواشف الضوئية بدمج القدرة الضوئية عبر جميع الأطوال الموجية، مما يحسب الضرب الداخلي التناظري (درجة التشابه) لجميع الـ $N$ من الكتل بالتوازي.
• سلسلة المعالجة (Pipeline):
  1. ترميز الاستعلام: يقوم المعالج الرسومي (GPU) بحساب مخطط استعلام (Sketch) ذي أبعاد $d$، ويقوم بتشفيره على مشط ليزر متعدد الأطوال الموجية (WDM) عبر محولات رقمية-تناظرية (DACs) ومعدلات.
  2. البث الضوئي: يقوم مقسم ضوئي سلبي $1 \times N$ ببث الاستعلام إلى جميع بنوك أوزان الـ MRR.
  3. حساب التشابه المتوازي: يحتوي كل مسار على $d$ من الـ MRRs المضبوطة على أطوال موجية محددة. تعمل هذه الـ MRRs كأوزان ($w_{n,j}$). يقوم الكاشف الضوئي بجمع الإشارات الموزونة، مما يعطي $I_n \propto \sum w_{n,j} s_j$ (الضرب الداخلي) لجميع الـ $N$ من الكتل في وقت واحد.
  4. الاختيار الإلكتروني: يتم تحويل التيارات التناظرية الـ $N$ إلى إشارات رقمية، ويقوم مقارن "أفضل $k$" رقمي باختيار الفهارس الأفضل.
  5. الجلب الانتقائي: يتم جلب أفضل $k$ من كتل الـ KV فقط من تخزين HBM/Flash، مما يقلل حركة مرور الذاكرة بشكل كبير.

الميزات العتادية الرئيسية

• المنصة: نيبات الليثيوم رقيقة الغشاء (TFLN)، والتي تم اختيارها بسبب تأثيرها الكهرو-بصري (Pockels) القوي، مما يسمح بضبط سريع ومنخفض الطاقة.
• الأوزان الموقعة: يستخدم النظام كشفاً ضوئياً متوازناً (عبر المنفذ النافذ ناقص المنفذ المفقود) لترميز الأوزان الموقعة ($[-1, +1]$) مباشرة، مما يتجنب الحاجة إلى عمليات إسقاط ReLU أو الترميز المنقسم.
• الأوزان شبه المستقرة: يتم تحديث أوزان الـ MRR (التواقيع) فقط عند اكتمال كتلة جديدة من التوكنات (كل $B$ من التوكنات)، مما يقلل من استهلاك الطاقة الديناميكي.

3. المساهمات الرئيسية

1. الرؤية الهيكلية: تحديد أن اختيار كتلة الـ KV هو بحث عن التشابه مقيد بالذاكرة، وهو ما يناسب تماماً بنيات البث والوزن الضوئية، مما يسمح بالقضاء على مسح الذاكرة من رتبة $O(N)$.
2. النمذجة المدركة للعتاد: تطوير نموذج شامل للاختلالات (الانجراف الحراري، التكميم، التداخل، الضوضاء) وإثبات أن النظام يحافظ على دقة عالية حتى مع وجود عيوب عتادية واقعية (دقة 4-6 بت، وانجراف حراري 30 بيكومتر).
3. تحليل رؤوس الاسترجاع: كشف تحليل رؤوس الانتباه في نماذج Qwen2.5-7B و Qwen3-8B أن أكثر من 90% من رؤوس الانتباه تعمل كـ "رؤوس استرجاع" عند السياقات الطويلة، مما يؤكد الحاجة إلى المسرع الضوئي لغالبية آلية الانتباه في النموذج.
4. التحقق الشامل: التحقق من النظام باستخدام اختبارات "الإبرة في كومة القش" (NIAH)، مما أثبت أن اختيار الكتل المتأثر بالـ MRR يحافظ على دقة الانتباه الكامل مع تقليل حركة مرور البيانات.

4. النتائج الرئيسية

• الدقة:
  • حقق دقة 100% في اختبار NIAH من 4K إلى 64K توكن (ضمن نافذة النموذج الأصلية) باستخدام $k=32$ من الكتل.
  • حتى في ظل الظروف العتادية المتشائمة (تكميم 4-بت، وانجراف حراري 30 بيكومتر)، ظلت الدقة مطابقة للانتباه الكامل.
  • عند 128K، عُزي تراجع الدقة إلى حدود النموذج الأساسي، وليس إلى آلية الاختيار.
• تقليل حركة مرور البيانات:
  • عند سياق 64K: تقليل بمقدار 16 ضعفاً في حركة مرور ذاكرة KV cache.
  • عند سياق 128K: تقليل بمقدار 32 ضعفاً.
  • عند 1M توكن: يُتوقع تقليل بمقدار ~244 ضعفاً.
• كفاءة الطاقة:
  • يستهلك PRISM حوالي 2.3 نانو جول لكل عملية اختيار استعلام.
  • مقارنة بمسح كامل من معالج H100 GPU (~16.3 ميكرو جول)، يوفر PRISM ميزة طاقة تبلغ ~7,000 ضعف عند سياق 128K.
  • حتى مقارنة بـ GPU يقوم بمسح التواقيع المضغوطة، يحافظ PRISM على ميزة تبلغ ~5,000 ضعف.
• زمن الانتقال (Latency):
  • زمن انتقال التقييم الضوئي هو ~9 نانو ثانية.
  • هذا أسرع بـ ~500 ضعف من المسح الكامل بواسطة GPU (~5 ميكرو ثانية).
  • نقطة التقاطع حيث يصبح PRISM مفيداً مقارنة بالأساليب الإلكترونية هي عند ~4K توكن.

5. الأهمية والآفاق المستقبلية

• تحول في النموذج المعرفي: ينقل PRISM التركيز من "الحوسبة الأسرع" إلى "الوصول الأذكى للذاكرة"، مما يحل قانون التوسع الأساسي لاستدلال الـ LLM (قراءة الذاكرة من رتبة $O(N)$) بجعلها من رتبة $O(1)$ عبر الفيزياء.
• القابلية للتوسع: تتوسع البنية بشكل إيجابي لتصل إلى سياقات بمليون توكن. وبينما تنمو قيود عرض النطاق الإلكتروني خطياً، فإن تكلفة النهج الضوئي تظل ثابتة لكل استعلام، مما يجعلها المسار الوحيد الممكن لاستدلال السياق فائق الطول.
• التكامل: يوضح البحث مساراً لدمج PRISM كرقاقة (chiplet) مترافق أو بطاقة PCIe، بحيث يعمل جنباً إلى جنب مع نظام ICMS من NVIDIA (المحول الذكي للاتصال والذاكرة) للجمع بين سعة التخزين الهائلة والاختيار الضوئي فائق السرعة.
• القدرة على التعميم: يمكن تطبيق نموذج "البحث بالبث الضوئي" لما وراء النماذج اللغوية الكبيرة، ليشمل أي عمل من أعمال البحث عن التشابه (مثل أنظمة التوصية، وقواعد بيانات المتجهات) حيث يجب مقارنة استعلام ما مع مجموعة كبيرة من المتجهات التي تتغير ببطء.