AI Infrastructure

توزیع بار حالت‌دار: مدیریت وابستگی کش KV در خوشه‌های توزیع‌شده LLM

استدلال توزیع‌شده مدل‌های زبانی بزرگ (LLM) از یک وظیفه پردازش دسته‌ای ایستا به یک عمل پویا و حالت‌دار تبدیل شده است. در حالی که سرورهای وب سنتی درخواست‌ها را به عنوان واحدهای مستقل در نظر می‌گیرند، LLMهای مدرن به شدت به کش Key-Value (KV) برای حفظ زمینه در چندین مرحله استدلال تکیه می‌کنند. هنگام استقرار در یک خوشه توزیع‌شده، جایگیری این درخواست‌ها حیاتی است. اگر درخواست بعدی برای مکالمه یک کاربر به یک گره GPU دیگر هدایت شود، مدل باید یا کل کش KV را از نو محاسبه کند یا به دلیل هزینه انتقال داده با تأخیرهای شدید روبرو شود. اینجاست که توزیع بار حالت‌دار به یک جزء اساسی زیرساخت هوش مصنوعی تبدیل می‌شود.

مشکل مسیریابی بدون حالت

توزیع‌کننده‌های بار استاندارد، مانند NGINX یا HAProxy، معمولاً بر اساس مبنای بدون حالت عمل می‌کنند و از الگوریتم‌هایی مانند Round Robin یا Least Connections برای توزیع ترافیک استفاده می‌کنند. در زمینه LLMها، این رویکرد بهینه نیست. کش KV به عنوان یک ردپای حافظه عمل می‌کند که با طول دنباله رشد می‌کند. جابه‌جایی این حالت بین GPUها پرهزینه است. اگر «وابستگی» (Affinity) را از دست بدهیم، عملکرد را از دست می‌دهیم.

یک برنامه چت چندگانه را در نظر بگیرید. پرامپت اول («برایم داستانی درباره یک گربه بگو») روی گره GPU A پردازش می‌شود. کش KV برای این پیشوند در حافظه HBM گره A ذخیره می‌شود. کاربر پاسخ می‌دهد («حالا آن را خنده‌دار کن»). اگر یک توزیع‌کننده بار بدون حالت این درخواست دوم را به گره GPU B هدایت کند، گره B هیچ آگاهی از زمینه قبلی ندارد. باید یا:

  1. کش KV را از گره A از طریق شبکه (PCIe/NVLink یا InfiniBand) دریافت کند، که محدودیت‌های پهنای باند قابل توجهی ایجاد می‌کند.
  2. فاز پیش‌پر (prefill) را برای کل زمینه قبلی محاسبه کند، که منجر به تکرار کار محاسباتی می‌شود.

پیاده‌سازی چسبندگی نشست با وابستگی KV

برای حل این مشکل، ما وابستگی کش KV را پیاده‌سازی می‌کنیم. این استراتژی شامل حفظ نگاری بین نشست‌های کاربر (یا پیشوندهای منحصربه‌فرد درخواست) و کارگرهای GPU خاص است. توزیع‌کننده بار قبل از مسیریابی یک درخواست جدید به این نقشه مراجعه می‌کند. اگر یک نشست قبلاً به یک کارگر تخصیص داده شده باشد، درخواست به آنجا هدایت می‌شود. فقط زمانی که یک کارگر اشباع شده یا در دسترس نباشد، توزیع‌کننده نشست را به یک گره جدید منتقل می‌کند و هزینه مهاجرت حالت یا محاسبه مجدد را به عنوان یک مصالحه ضروری می‌پذیرد.

ملاحظات معماری

یک توزیع‌کننده بار مؤثر به یک مخزن متمرکز یا توزیع‌شده (مانند Redis یا etcd) برای حفظ نقشه وابستگی نیاز دارد. تصمیمات کلیدی طراحی شامل موارد زیر است:

  • تعریف کلید: استفاده از شناسه نشست ساده اما سفت و سخت است. استفاده از هش پیشوند پرامپت انعطاف‌پذیری بهتری برای کش معنایی ارائه می‌دهد اما هزینه محاسباتی را در دروازه افزایش می‌دهد.
  • بررسی سلامت: توزیع‌کننده باید از مصرف حافظه GPU و دما نظارت کند. اگر یک کارگر که یک نشست «داغ» را نگه می‌دارد به محدودیت حافظه نزدیک شود، سیستم باید به طور پیش‌دستانه نشست را اخراج کند یا آن را به طور نرم قبل از وقوع خطای OOM (خارج از حافظه) منتقل کند.
  • ریزش بار (Load Shedding): وقتی همه گره‌های دارای وابستگی بار اضافی دارند، سیستم باید نشست‌های جدید را به گره‌های کمتر استفاده‌شده اولویت‌بندی کند به جای صف شدن نامحدود پشت بارهای موجود با وابستگی بالا.

نمونه کد: منطق مسیریابی مبتنی بر وابستگی

در زیر یک پیاده‌سازی مفهومی پایتون از یک توزیع‌کننده بار حالت‌دار آمده است که وابستگی را اولویت می‌دهد و در عین حال محدودیت‌های بار را رعایت می‌کند.


import random
from collections import defaultdict

class StatefulLLMLoadBalancer:
    def __init__(self, max_workers_per_session=1):
        # Maps session_id -> worker_id
        self.affinity_map = {}
        # Maps worker_id -> list of active session_ids
        self.worker_sessions = defaultdict(list)
        # Simulated worker load (e.g., number of active requests or memory usage)
        self.worker_load = defaultdict(int)

    def assign_worker(self, session_id: str, current_loads: dict) -> str:
        """
        Assigns a worker for a given session ID.
        Prioritizes existing affinity if the worker is under the load threshold.
        """
        threshold = 10  # Maximum concurrent sessions per worker before migration

        # 1. Check if we have an existing affinity
        if session_id in self.affinity_map:
            current_worker = self.affinity_map[session_id]
            
            # If the current worker is not overloaded, stick to it
            if current_loads.get(current_worker, 0) < threshold:
                return current_worker
            
            # If overloaded, we must migrate. Remove from old worker context.
            print(f"Warning: Worker {current_worker} overloaded. Migrating session {session_id}.")
            self.worker_sessions[current_worker].remove(session_id)

        # 2. Select a new worker
        # Strategy: Least Loaded Worker among those not at capacity
        available_workers = [
            w for w, load in current_loads.items() 
            if load < threshold
        ]
        
        if not available_workers:
            raise Exception("System Capacity Full: No available workers under threshold.")

        # Pick the worker with the lowest current load
        new_worker = min(available_workers, key=lambda w: current_loads[w])

        # 3. Update Affinity Maps
        self.affinity_map[session_id] = new_worker
        if session_id not in self.worker_sessions[new_worker]:
            self.worker_sessions[new_worker].append(session_id)
        
        return new_worker

    def simulate_batch(self, sessions: list[str], initial_loads: dict):
        """Simulates routing a batch of requests."""
        routing_decisions = []
        for session in sessions:
            worker = self.assign_worker(session, initial_loads)
            initial_loads[worker] += 1
            routing_decisions.append((session, worker))
        return routing_decisions

# Usage Example
lb = StatefulLLMLoadBalancer()
# Simulate initial state: Worker 1 has 9 sessions, Worker 2 has 2
current_loads = {"worker_1": 9, "worker_2": 2}

# Session "abc" was previously assigned to worker_1 (simulated by pre-populating map)
lb.affinity_map["abc"] = "worker_1"
lb.worker_sessions["worker_1"].append("abc")

# New batch of incoming requests
new_requests = ["abc", "def", "ghi"]

print("Routing Decisions:")
for session, worker in lb.simulate_batch(new_requests, current_loads):
    print(f"Session: {session:5s} -> Routed to: {worker}")

# Output Explanation:
# Session 'abc' should route to worker_1 if load < 10. 
# Since worker_1 load is 9, it stays there. Load becomes 10.
# Session 'def' has no affinity, goes to least loaded (worker_2). Load becomes 3.
# Session 'ghi' has no affinity, goes to least loaded (worker_2). Load becomes 4.

استراتژی‌های پیشرفته: وابستگی معنایی و مهاجرت پیش‌دستانه

برای استقرارهای مقیاس بالا، چسبندگی ساده شناسه نشست ممکن است کافی نباشد. سیستم‌های پیشرفته از وابستگی معنایی استفاده می‌کنند، جایی که توزیع‌کننده بار توکن‌های اولیه پرامپت را هش می‌کند. اگر دو کاربر مختلف مکالمه را با یک پرامپت سیستم یکسان شروع کنند (مثلاً «تو یک دستیار برنامه‌نویسی مفید هستی»), کش‌های KV آن‌ها برای آن پیشوند می‌توانند به اشتراک گذاشته شوند یا روی یک گره نگه داشته شوند. این امر به ویژه برای برنامه‌های RAG (تولید تقویت‌شده با بازیابی) مؤثر است که در آن‌ها پنجره‌های زمینه بزرگ و ایستا هستند.

علاوه بر این، پیاده‌سازی مهاجرت پیش‌دستانه حیاتی است. به جای اینکه منتظر کرش شدن یک کارگر به دلیل OOM بمانیم، توزیع‌کننده باید فشار حافظه را نظارت کند. وقتی یک کارگر از ۸۵٪ استفاده حافظه فراتر رود، می‌تواند به توزیع‌کننده سیگنال بدهد تا نشست‌های جدید را در جای دیگر «شروع سرد» کند و در صورت امکان، کش‌های KV نشست‌های کم‌فعال‌تر را سریال کرده و به ذخیره‌سازی کندتر با ظرفیت بزرگ‌تر (مانند RAM پردازنده مرکزی یا SSDهای NVMe) منتقل کند تا حافظه GPU را برای درخواست‌های با اولویت بالا آزاد کند.

نتیجه‌گیری

توزیع بار حالت‌دار دیگر یک ویژگی مطلوب نیست، بلکه یک الزام اصلی برای زیرساخت LLM کارآمد است. با مدیریت وابستگی کش KV و چسبندگی نشست، می‌توانیم زمان تا اولین توکن (TTFT) را به طور قابل توجهی کاهش دهیم و بهره‌برداری کلی از خوشه را بهبود بخشیم. کلید کار در تعادل بین هزینه مهاجرت حالت و مزیت محلی بودن کش است. با ادامه رشد LLMها در اندازه و طول پنجره زمینه، پیچیدگی این الگوریتم‌های مسیریابی فقط افزایش خواهد یافت و ما را از نشست‌های چسبنده ساده به سمت ارکستراسیون هوشمند و آگاه از معنای منابع هدایت می‌کند. برای توسعه‌دهندگانی که بک‌اند هوش مصنوعی می‌سازند، یکپارچه‌سازی منطق حالت‌دار در لایه توزیع بار در مراحل اولیه برای مقیاس‌پذیری مؤثر بدون تحمیل هزینه‌های محاسباتی غیرقابل قبول، حیاتی است.

Share: