AI Infrastructure

Durumlu Yük Dengeleme: Dağıtık LLM Kümelemede KV Önbellek Bağlılığının Yönetimi

Dağıtık Büyük Dil Modeli (LLM) çıkarımı, statik bir toplu işleme görevinden dinamik ve durumlu bir operasyona dönüştü. Geleneksel web sunucuları istekleri bağımsız birimler olarak ele alırken, modern LLM'ler bağlamı birden fazla çıkarım adımı boyunca korumak için Key-Value (KV) önbelleğine büyük ölçüde güveniyor. Dağıtık bir kümede dağıtıldığında, bu isteklerin yerleşimi kritik hale geliyor. Bir kullanıcının sohbetine ait sonraki bir istek farklı bir GPU düğümüne yönlendirilirse, modelin tüm KV önbelleğini sıfırdan yeniden hesaplaması veya veri aktarım maliyeti nedeniyle ciddi gecikme cezalarıyla karşılaşması gerekir. İşte durumlu yük dengeleme, yapay zeka altyapısının temel bir bileşeni olarak burada devreye giriyor.

Durumsuz Yönlendirmedeki Sorun

NGINX veya HAProxy gibi standart yük dengeleyiciler, trafiği dağıtmak için Round Robin veya Least Connections gibi algoritmaları kullanarak genellikle durumsuz bir şekilde çalışır. LLM'ler bağlamında bu yaklaşım en iyisi değildir. KV önbellek, dizin uzunluğuyla birlikte büyüyen bir bellek izi görevi görür. Bu durumu GPU'lar arasında taşımak pahalıdır. "Bağlılığı" (affinity) kaybedersek, performansı da kaybederiz.

Çok turlu bir sohbet uygulamasını düşünün. İlk istem ("Bana bir kedi hakkında hikaye anlat") GPU Düğüm A üzerinde işlenir. Bu önekin KV önbelleği, Düğüm A'nın HBM belleğinde saklanır. Kullanıcı yanıt verir ("Şimdi onu komik yap"). Durumsuz bir yük dengeleyici bu ikinci isteği GPU Düğüm B'ye yönlendirirse, Düğüm B önceki bağlam hakkında hiçbir bilgiye sahip değildir. Ya şunu yapmalıdır:

  1. KV önbelleğini ağ üzerinden (PCIe/NVLink veya InfiniBand) Düğüm A'dan getirmelidir; bu, önemli bant genişliği kısıtlamaları getirir.
  2. Tüm önceki bağlam için prefill aşamasını yeniden hesaplamalıdır; bu da hesaplama işini yineler.

KV Bağlılığı ile Oturum Yapışkanlığının Uygulanması

Bunu çözmek için KV Önbellek Bağlılığı (KV Cache Affinity) uygularız. Bu strateji, kullanıcı oturumları (veya benzersiz istem önekleri) ile belirli GPU işçileri arasındaki eşlemeyi korumayı içerir. Yük dengeleyici, yeni bir isteği yönlendirmeden önce bu haritaya danışır. Bir oturum zaten bir işçiye atanmışsa, istek oraya yönlendirilir. Yalnızca bir işçi doymuş veya kullanılamaz durumda olduğunda, dengeleyici oturumu yeni bir düğüme taşır ve durum taşınması veya yeniden hesaplama maliyetini gerekli bir ödünleşim olarak kabul eder.

Mimari Düşünceler

Etkili bir durumlu dengeleyici, bağlılık haritasını sürdürmek için merkezi veya dağıtık bir depolama (Redis veya etcd gibi) gerektirir. Temel tasarım kararları şunları içerir:

  • Anahtar Tanımı: Oturum kimliğini kullanmak doğrudan ama katıdır. İstem önekinin bir hash'ini kullanmak, anlamsal önbellekleme için daha iyi esneklik sunar ancak kapıdaki hesaplama maliyetini artırır.
  • Sağlık Kontrolleri: Dengeleyici, GPU bellek kullanımını ve sıcaklığını izlemelidir. "Sıcak" bir oturumu tutan bir işçi bellek sınırlarına yaklaşıyorsa, sistem, bir OOM (Bellek Dışı) hatası oluşmadan önce oturumu proaktif olarak atmalı veya zarifçe taşımak zorundadır.
  • Yük Düşürme (Load Shedding): Bağlılığı olan tüm düğümler aşırı yüklendiğinde, sistem, mevcut yüksek bağlılık yüklerinin arkasında süresiz kuyrukta beklemek yerine, yeni oturumlara daha az kullanılan düğümleri önceliklendirmelidir.

Kod Örneği: Bağlılığa Dayalı Yönlendirme Mantığı

Aşağıda, bağlılığı önceliklendirirken yük sınırlarına saygı gösteren durumlu bir yük dengeleyicisinin kavramsal bir Python uygulaması yer almaktadır.


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.

Gelişmiş Stratejiler: Anlamsal Bağlılık ve Önleyici Taşıma

Yüksek ölçekli dağıtımlar için basit oturum kimliği yapışkanlığı yeterli olmayabilir. Gelişmiş sistemler, yük dengeleyicinin istemin ilk birkaç token'ını hash'lediği anlamsal bağlılığı (semantic affinity) kullanır. İki farklı kullanıcı aynı sistem istemiyle (ör. "Sen yardımsever bir kodlama asistanısın") bir sohbet başlatırsa, bu önek için KV önbellekleri paylaşılabilir veya aynı düğümde tutulabilir. Bu, bağlam pencerelerinin büyük ve statik olduğu RAG (Geri Getirme Destekli Üretim) uygulamaları için özellikle etkilidir.

Ayrıca, önleyici taşımayı (preemptive migration) uygulamak kritiktir. Bir işçinin OOM nedeniyle çökmesini beklemek yerine, dengeleyici bellek baskısını izlemelidir. Bir işçi %85 bellek kullanımını aştığında, dengeleyiciye yeni oturumları başka yerde "soğuk başlatması" için sinyal verebilir ve mümkünse, daha az aktif oturumların KV önbelleklerini serileştirip daha yavaş, daha büyük kapasiteli depolamaya (CPU RAM veya NVMe SSD'ler gibi) aktararak yüksek öncelikli istekler için GPU belleğini serbest bırakabilir.

Sonuç

Durumlu yük dengeleme artık istenirse güzel olur özelliği değil, verimli LLM altyapısı için temel bir gerekliliktir. KV önbellek bağlılığını ve oturum yapışkanlığını yöneterek, ilk token'a kadar geçen süreyi (TTFT) önemli ölçüde azaltabilir ve genel küme kullanımını iyileştirebiliriz. Anahtar nokta, durum taşınma maliyetini önbellek yerelliğinin faydasıyla dengelemektir. LLM'ler boyut ve bağlam pencere uzunluğu olarak büyümeye devam ettikçe, bu yönlendirme algoritmalarının karmaşıklığı artacak ve bizi basit yapışkan oturumlardan akıllı, anlamsal farkında kaynak orkestrasyonuna doğru taşıyacaktır. Yapay zeka arka uçları geliştiren geliştiriciler için, durumlu mantığı yük dengeleme katmanına erken entegre etmek, aşırı hesaplama maliyetleri katlanmadan etkili ölçeklenmek için kritiktir.

Share: