‏FlashMemory DS-V4: نگه‌داشتن فقط ۱۰ تا ۱۵ درصد کش روی GPU

این مخزن یه retriever سبک به اسم FlashMemory DS-V4 معرفی می‌کنه که برای مدل DeepSeek-V4 ساخته شده و هدفش کم کردن مصرف حافظهٔ KV-cache هست. مدل‌های زبانی بزرگ موقع تولید متن، کلید/مقدار توکن‌های قبلی رو تو کشی نگه می‌دارن که تو کانتکست‌های طولانی خیلی بزرگ و سنگین می‌شه. ایده اینه که لازم نیست همهٔ این کش همیشه روی GPU بمونه.

منطق کار اینه: با گرفتن hidden state یه توکن در حال decode، این retriever پیش‌بینی می‌کنه که ۶۴ توکن بعدی به کدوم تکه‌های (chunk) کشِ فشردهٔ CSA توجه می‌کنن. فقط تکه‌های پرامتیاز روی دستگاه می‌مونن و بقیه به CPU یا دیسک منتقل می‌شن. به همین خاطر اسمش «lookahead» هست، چون جلوتر رو نگاه می‌کنه تا تصمیم بگیره چی نگه داره.

استفاده‌اش هم ساده‌ست: مدل رو از روی checkpoint بار می‌زنی و بهش hidden state، کلیدهای فشرده و موقعیت توکن‌ها رو می‌دی، اون هم یه ماسک نگه‌داری برمی‌گردونه:

model = FlashMemoryRetriever.from_checkpoint(
    "weights/flashmemory_ds_v4.safetensors", device="cuda"
)
scores = model.ensemble(hidden, comp_k, positions, mode="max")
keep = model.select_topk(hidden, comp_k, positions, top_k=512)

از نظر معماری، هر تکهٔ کش به‌صورت ۱۳۲ بایت uint8 ذخیره می‌شه (۱۲۸ بایت کلیدِ کوانتیزه‌شده با float8_e4m3 به‌علاوهٔ ۴ بایت مقیاس). امتیاز هر تکه از یه مسیر محاسباتی شامل پروجکشن کوئریِ LoRA، نرمال‌سازی RMSNorm، اعمال RoPE از نوع YaRN و تبدیل Hadamard می‌گذره و آخرش با یه sigmoid بین صفر و یک درمیاد. نکتهٔ جالب اینه که checkpoint سه لایهٔ CSA مستقل (l10، l12، l20) رو نگه می‌داره و امتیازهاشون به‌صورت max یا mean ترکیب می‌شن تا یه تصمیم نهاییِ نگه‌دار/بنداز بسازن.

طبق ارزیابی‌های آورده‌شده، FlashMemory روی کارهای استدلالیِ کانتکست‌بلند مثل RULER، LongMemEval و LongBench V2 هم‌سطح یا کمی بهتر از حالت توجه کامل عمل می‌کنه، در حالی که فقط حدود ۱۰ تا ۱۵ درصد کش رو روی دستگاه نگه می‌داره و تا حدود ۹۰ درصد حافظه صرفه‌جویی می‌شه. فقط کارهای سوزن‌در‌انبارکاهِ دقیق (مثل MRCR) به یه سازوکار threshold-fallback اضافه تو لایهٔ سرویس نیاز دارن که جزو نسخهٔ مستقل نیست.

یه نکتهٔ صادقانه هم تو مخزن هست: کد منتشرشده فقط یه نمونهٔ آموزشی torch-only برای نشون دادن جریان کنترل decode هست، نه یه DeepSeek-V4 قابل‌اجرا. موتور واقعیِ جابه‌جایی کش بین CPU و GPU و فریم‌ورک داخلی sglang منتشر نشده و قابل‌انتشار هم نیست. وزن‌های آموزش‌دیده (حدود ۵۱۰ مگابایت) روی Hugging Face موجوده و مجوز پروژه MIT هست.

نکات کلیدی:

• یه retriever که حدس می‌زنه ۶۴ توکن بعدی به کدوم تکه‌های KV-cache توجه می‌کنن
• فقط تکه‌های پرامتیاز روی GPU می‌مونن، بقیه به CPU/دیسک منتقل می‌شن
• با نگه‌داشتن فقط ۱۰ تا ۱۵ درصد کش، هم‌سطح یا بهتر از توجه کامل تو کانتکست بلند
• checkpoint سه لایهٔ CSA رو ترکیب می‌کنه و کلیدها با fp8 کوانتیزه شدن
• کد منتشرشده فقط آموزشیه؛ موتور واقعی جابه‌جایی کش منتشر نشده. مجوز MIT، وزن‌ها روی Hugging Face