MISA: Mixture of Indexer Sparse Attention for Long-Context LLM Inference

ملخص تقني: MISA (خليط من رؤوس الفهرسة للاهتمام المتناثر - Mixture of Indexer Sparse Attention)

بيان المشكلة
مع توسع النماذج اللغوية الكبيرة (LLMs) لمعالجة سياقات تصل إلى مئات الآلاف أو ملايين الرموز (tokens)، تصبح آليات الاهتمام الكثيفة (dense attention) مكلفة حاسوبياً بشكل باهظ. وبينما ظهرت تقنيات الاهتمام المتناثر (sparse attention) للتخفيف من ذلك، يمثل "اهتمام ديب سيك المتناثر" (DeepSeek Sparse Attention - DSA) حالياً أحدث ما توصل إليه العلم في مجال الاهتمام المتناثر دقيق التفاصيل أثناء الاستدلال. يستخدم DSA "فهرساً" (indexer) متعلماً يقوم بتقييم كل رمز بادئة (prefix token) باستخدام مجموعة كبيرة من الرؤوس $H_I$ (على سبيل المثال، 64 رأساً في DeepSeek-V3.2) لاختيار أفضل $k$ من الرموز لآلية الاهتمام الرئيسية.

وعلى الرغم من أن تصميم تعدد الرؤوس في فهرس DSA يعد أمراً حيوياً للتعبير (expressiveness) — مما يسمح للرؤوس المختلفة بالتخصص في أنماط صلة متميزة (مثل الحداثة، أو البنية النحوية، أو الدلالة) — إلا أنه يخلق عنق زجاجة كبيراً. فالتكلفة الحسابية لتقييم كل رمز باستخدام جميع الرؤوس $H_I$ تتناسب طردياً مع $O(H_I L)$، مما يجعل الفهرس هو التكلفة المهيمنة أثناء الاستدلال في السياقات الطويلة. إن التحسينات الحالية مثل HISA (البحث الهرمي من الكتلة إلى الرمز) وIndexCache (إعادة استخدام الفهرس عبر الطبقات) تقلل التكاليف من خلال العمل على محور الرمز أو محور الطبقة، لكنها تظل تتطلب بقاء جميع رؤوس $H_I$ نشطة داخل النواة (kernel)، مما يعني أنها تفشل في معالجة عنق زجاجة محور الرؤوس بشكل مباشر.

المنهجية: MISA
يقترح المؤلفون MISA (خليط من رؤوس الفهرسة للاهتمام المتناثر)، وهو بديل جاهز للتركيب (drop-in replacement) لفهرس DSA، يعامل رؤوس الفهرس $H_I$ كمجموعة من "خليط الخبراء" (Mixture-of-Experts - MoE). بدلاً من تفعيل جميع الرؤوس لكل استعلام، يقوم MISA بتوجيه مجموعة فرعية $h \ll H_I$ من الرؤوس النشطة المعتمدة على الاستعلام.

تتضمن الآلية الجوهرية مرحلتين:

1. الموجه المجمع للكتل خفيف الوزن (Lightweight Block-Pooled Router): يتم تقسيم البادئة إلى $M = \lceil L/B \rceil$ من الكتل. يقوم الموجه بحساب التقاربات لكل رأس باستخدام إحصائيات مجمعة على مستوى الكتلة (متوسط مفاتيح الفهرسة) بدلاً من الرموز الفردية. يقوم بتجميع هذه الدرجات لتقدير أهمية كل رأس للاستعلام الحالي واختيار أفضل $h$ من الرؤوس. تعمل هذه الخطوة على $M \ll L$ من المفاتيح، مما ينتج عنه عبء حسابي ضئيل قدره $O(H_I M)$.
2. تقييم الرموز المتناثر (Sparse Token Scoring): تقوم الرؤوس $h$ المختارة فقط بإجراء عملية التقييم الثقيلة على مستوى الرمز للمقدمة الكاملة. يتم بعد ذلك اختيار أفضل $k$ من الرموز بناءً على هذه الدرجات المختزلة.

يقلل هذا النهج من تعقيد الفهرس لكل استعلام من $O(H_I L)$ إلى $O(hL + H_I M)$. وبما أن $M \ll L$ و $h \ll H_I$، فإن التوفير في الحوسبة يكون كبيراً مع الحفاظ على التنوع الكامل لمجموعة الرؤوس (تظل جميع الرؤوس متاحة؛ الموجه ببساطة يختار أيها سيستشير).

النسخة الهرمية (MISA†)
لسد الفجوة بين الكفاءة والاختيار الدقيق لنموذج DSA الكثيف، قدم المؤلفون MISA†، وهو مسار من "الخشن إلى الدقيق" (coarse-to-fine):

• المرحلة الخشنة (Coarse Stage): يختار المسار الموجه عبر MoE مجموعة مرشحة موسعة من $k'$ من الرموز (على سبيل المثال، $k'=8192$).
• المرحلة الدقيقة (Fine Stage): يقوم فهرس DSA الأصلي (باستخدام جميع رؤوس $H_I$) بإعادة تقييم هذه المجموعة المرشحة فقط لاستعادة أفضل $k$ الرموز.
  يسمح هذا لـ MISA† باستعادة اختيارات الفهرس الكثيف بدقة شبه تامة مع الاستفادة من الحوسبة المنخفضة للمرحلة الخشنة.

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

1. توجيه محور الرؤوس (Head-Axis Routing): حددت الورقة أن نواتج (products) الرؤوس والرموز لكل رمز هي التكلفة المهيمنة في DSA، وقدمت توجيه محور الرؤوس كأفق كفاءة مكمل لمحاور الهرمية في محور الرموز الموجودة حالياً.
2. بنية MISA: فهرس موجه بأسلوب MoE جاهز للتركيب، يقوم بتفعيل $h$ فقط من الرؤوس لكل استعلام عبر موجه مجمع للكتل خفيف الوزن، مما يقلل التعقيد دون التضحية بتنوع مجموعة الرؤوس.
3. امتداد MISA†: نسخة هرمية تستخدم المسار الموجه لتقليص مساحة البحث، يليه إعادة ترتيب كامل لـ DSA، مما يحقق استعادة شبه مثالية لاختيارات الفهرس الكثيف.
4. الكفاءة والأداء: يحقق المنهج تسريعاً كبيراً دون الحاجة إلى تدريب إضافي.

النتائج التجريبية
تم التقييم على DeepSeek-V3.2 ($H_I=64$) و GLM-5 ($H_I=32$) دون تدريب إضافي:

• LongBench: يطابق MISA (مع $h=8$) فهرس DSA الكثيف ضمن 0.20 نقطة في المتوسط على DeepSeek-V3.2 ويتفوق عليه في GLM-5. كما يقلص MISA† الفجوة لتصبح ضمن 0.1 نقطة على DeepSeek-V3.2. يتفوق كلاهما على نماذج Block-Sparse و HISA كمرجع أساسي.
• الإبرة في كومة القش (Needle-in-a-Haystack - NIAH): يحافظ MISA و MISA† على خرائط حرارية خضراء بالكامل للاسترجاع حتى سياق 128K، مما يطابق أداء DSA الكثيف. في المقابل، يظهر Block-Sparse ثقوباً في الدقة، ويظهر HISA تدهورات طفيفة عند الأعماق القصوى.
• استعادة الرموز: يستعيد MISA† أكثر من 92% من الرموز التي اختارها DSA الكامل لكل طبقة (مقاسة بـ IoU).
• التسريع: تحقق نسخة نواة TileLang تسريعاً قدره 3.82× تقريباً مقارنة بنواة DSA الأصلية على وحدة معالجة رسومات واحدة من نوع NVIDIA H200 للسياقات الطويلة.

الأهمية والادعاءات
تدعي الورقة أن توجيه محور رأس الفهرس (indexer-head-axis routing) هو محور كفاءة عملي ومكمل للاهتمام المتناثر دقيق التفاصيل. من خلال فصل قرار التوجيه (أي الرؤوس المستخدمة) عن اختيار الرموز (أي الرموز التي يتم الانتباه إليها)، يثبت MISA أنه يمكن تقليل التكلفة الحسابية للفهرس بشكل جذري (بمقدار 8 أضعاف في DeepSeek-V3.2 و 4 أضعاف في GLM-5) مع الحفاظ على جودة الاسترجاع والقدرة التعبيرية للفهرس الكثيف الكامل. ويضع المؤلفون هذا كتحسين أساسي يمكن دمجه بشكل متعامد مع أي مخطط اهتمام متناثر موجود على مستوى الرمز.

القيود
يقر المؤلفون بأن تجارب السرعة لديهم تقيس زمن انتقال النواة (kernel latency) وليس زمن انتقال النموذج بالكامل، وأن الطريقة لا تقلل من حجم الوصول إلى الذاكرة (memory access volume) لذاكرة التخزين المؤقت (KV cache). بالإضافة إلى ذلك، تم الحصول على النتائج عن طريق إدخال MISA في نماذج مدربة مسبقاً دون تدريب مشترك للموجه، مما يشير إلى إمكانية تحقيق مكاسب إضافية من خلال التدريب.