본문으로 건너뛰기

[sglang] Mamba 호스트 캐시 메커니즘 최적화: 성능 향상과 메모리 관리 개선

PR 링크: sgl-project/sglang#21750 상태: Merged | 변경: +None / -None

들어가며

최근 LLM(거대 언어 모델) 분야에서는 Mamba와 같은 새로운 아키텍처가 등장하며 기존의 Transformer 모델과는 다른 방식으로 성능을 끌어올리고 있습니다. 특히 Mamba는 상태 공간 모델(State Space Model, SSM)을 기반으로 하여 긴 시퀀스 처리에서 뛰어난 효율성을 보여줍니다. 하지만 이러한 모델을 효율적으로 서빙하기 위해서는 메모리 관리 및 캐싱 전략이 매우 중요합니다.

이번 PR, "[HiMambaTree]: Optimize mamba host lock mechanism"은 sglang 프로젝트에서 Mamba 모델의 호스트 캐시 메커니즘을 최적화하여 메모리 사용량을 줄이고 추론 성능을 향상시키는 것을 목표로 합니다. 기존에는 Mamba 모델의 호스트 캐시 관리에 다소 비효율적인 부분이 존재했으며, 이 PR은 이를 개선하여 더 빠르고 효율적인 서빙을 가능하게 합니다.

본 글에서는 이 PR의 변경 사항을 상세히 분석하고, 왜 이러한 변경이 성능 향상에 기여하는지, 그리고 어떤 일반적인 교훈을 얻을 수 있는지 살펴보겠습니다.

코드 분석

이번 PR의 핵심 변경 사항은 python/sglang/srt/mem_cache/hi_mamba_radix_cache.py와 python/sglang/srt/mem_cache/mamba_radix_cache.py 파일에 집중되어 있습니다. 주요 변경 내용은 Mamba 모델의 호스트 캐시(host cache)에 대한 참조 카운터(reference counter) 관리와 보호(protection) 로직을 개선하는 것입니다.

1. mamba_radix_cache.py에서의 TreeNode 클래스 수정

가장 근본적인 변화는 Mamba 모델의 각 노드(node)가 호스트 캐시뿐만 아니라 호스트 Mamba 캐시(host mamba cache)에 대한 참조 카운터도 별도로 관리하도록 변경된 점입니다.

Before:

class TreeNode:
    # ...
    self.hit_count = 0
    self.host_ref_counter = 0
    # ...

After:

class TreeNode:
    # ...
    self.hit_count = 0
    self.host_ref_counter = 0
    self.host_mamba_ref_counter = 0
    # ...

    def protect_host(self):
        """Protect the host KV value from eviction."""
        self.host_ref_counter += 1

    def release_host(self):
        """Release the host KV value, allowing it to be evicted."""
        if self.host_ref_counter > 0:
            self.host_ref_counter -= 1
        else:
            raise RuntimeError("Host reference counter is already zero.")

    def protect_host_mamba(self):
        """Protect the host mamba value from eviction."""
        self.host_mamba_ref_counter += 1

    def release_host_mamba(self):
        """Release the host mamba value, allowing it to be evicted."""
        if self.host_mamba_ref_counter > 0:
            self.host_mamba_ref_counter -= 1
        else:
            raise RuntimeError("Host mamba reference counter is already zero.")

설명:

  • TreeNode 클래스에 host_mamba_ref_counter라는 새로운 필드가 추가되었습니다. 이는 호스트에 저장된 Mamba 관련 데이터(예: Mamba의 상태 벡터 등)에 대한 참조 횟수를 추적합니다.
  • protect_host_mamba()와 release_host_mamba() 메소드가 새로 추가되었습니다. 이 메소드들은 각각 호스트 Mamba 캐시를 보호하거나 해제하는 역할을 하며, host_mamba_ref_counter를 증가시키거나 감소시킵니다. 이는 기존의 host_ref_counter가 주로 KV 캐시와 같은 일반적인 호스트 메모리 할당을 관리하는 것과 별개로 작동합니다.
  • release_host() 메소드에서 host_ref_counter가 0 이하로 내려가는 경우 RuntimeError를 발생시키도록 하여 참조 카운터의 무결성을 강화했습니다. 이는 리뷰어 ispobock의 피드백을 반영한 것으로 보입니다.

2. hi_mamba_radix_cache.py에서의 로직 수정

HiMambaRadixCache 클래스 내의 여러 메소드에서 Mamba 호스트 캐시 참조 카운터(host_mamba_ref_counter)를 고려하도록 로직이 수정되었습니다.

_protect_host_node 및 _release_host_node 메소드:

이 메소드들은 노드를 호스트에 보호하거나 해제하는 핵심 로직을 담당합니다. Mamba 관련 보호/해제 로직이 추가되었습니다.

Before:

    def _protect_host_node(self, node: TreeNode):
        node.protect_host()
        self.evictable_full_host_leaves.discard(node)

    def _release_host_node(self, node: TreeNode):
        node.release_host()
        if node.host_ref_counter == 0:
            self._update_full_host_leaf_status(node)

After:

    def _protect_host_node(self, node: TreeNode, protect_mamba: bool = True):
        node.protect_host()
        self.evictable_full_host_leaves.discard(node)
        if protect_mamba:
            node.protect_host_mamba()
            if self.mamba_host_lru_list.in_list(node):
                self.mamba_host_lru_list.remove_node(node)

    def _release_host_node(self, node: TreeNode, release_mamba: bool = True):
        node.release_host()
        if release_mamba:
            node.release_host_mamba()
            if node.host_mamba_ref_counter == 0 and node.mamba_host_value is not None:
                if not self.mamba_host_lru_list.in_list(node):
                    self.mamba_host_lru_list.insert_mru(node)
        if node.host_ref_counter == 0 and node.host_mamba_ref_counter == 0:
            self._update_full_host_leaf_status(node)

설명:

  • _protect_host_node는 이제 protect_mamba 인자를 받아 Mamba 캐시 보호 여부를 결정합니다. protect_mamba=True일 경우 node.protect_host_mamba()를 호출하고, 해당 노드가 Mamba LRU 리스트에 있다면 제거합니다. 이는 Mamba 캐시가 보호되는 동안에는 LRU에서 제외되어 불필요한 evict 대상이 되지 않도록 합니다.
  • _release_host_node 역시 release_mamba 인자를 받습니다. release_mamba=True일 경우 node.release_host_mamba()를 호출합니다. Mamba 캐시 참조 카운터가 0이 되고 Mamba 값이 존재하는 경우, 해당 노드를 Mamba LRU 리스트에 MRU(Most Recently Used)로 삽입합니다. 이는 Mamba 캐시가 해제된 노드를 다시 사용할 수 있도록 준비하는 과정입니다.
  • _update_full_host_leaf_status 호출 조건이 node.host_ref_counter == 0에서 node.host_ref_counter == 0 and node.host_mamba_ref_counter == 0으로 변경되었습니다. 이는 노드의 호스트 KV 캐시와 호스트 Mamba 캐시가 모두 해제되었을 때만 해당 노드를 'evictable_full_host_leaves' 집합에서 업데이트하도록 하여, Mamba 캐시가 여전히 사용 중일 때는 해당 노드가 evict 대상에서 제외되도록 합니다.

_update_full_host_leaf_status 메소드:

evictable_full_host_leaves 집합에 노드를 추가하거나 제거하는 조건을 수정했습니다.

Before:

            or node.host_ref_counter > 0

After:

            or node.host_ref_counter > 0
            or node.host_mamba_ref_counter > 0

설명:

  • 노드가 'evictable_full_host_leaves' 집합에 포함되기 위한 조건으로 node.host_mamba_ref_counter > 0이 추가되었습니다. 즉, 호스트 Mamba 캐시가 사용 중인 노드는 더 이상 즉시 evict 가능한 상태로 간주되지 않습니다. 이는 Mamba 캐시의 무결성을 보장하고, Mamba 관련 데이터가 불필요하게 제거되는 것을 방지합니다.

_evict_host_leaf 및 _delete_tombstone_leaf 메소드:

호스트 리프 노드를 evict하거나 tombstone을 삭제할 때, 호스트 Mamba 캐시 참조 카운터가 0인지 확인하는 조건이 추가되었습니다.

Before (_evict_host_leaf):

        assert (
            node.host_ref_counter == 0
        ), f"in use, {node.id=} {node.host_ref_counter=}"

After (_evict_host_leaf):

        assert (
            node.host_ref_counter == 0
        ), f"host kv in use, {node.id=} {node.host_ref_counter=}"
        assert (
            node.host_mamba_ref_counter == 0
        ), f"host mamba in use, {node.id=} {node.host_mamba_ref_counter=}"

설명:

  • 호스트 리프 노드를 evict하기 전에 node.host_mamba_ref_counter == 0인지 확인하는 assert 문이 추가되었습니다. 이는 호스트 Mamba 캐시가 사용 중인 노드는 evict될 수 없음을 명확히 합니다.

evict_mamba_host 메소드:

호스트 Mamba 캐시를 evict하는 로직에서 Mamba 캐시 참조 카운터가 0인지 확인하는 조건이 추가되었습니다.

Before:

            if x.host_ref_counter > 0:
                x = x_next
                continue

After:

            # Leaf: evictable_full_host_leaves guarantees both counters == 0
            assert (
                x.host_ref_counter == 0
            ), f"evict host leaf: host_ref_counter != 0 with {x.id=} {x.host_ref_counter=}"
            assert (
                x.host_mamba_ref_counter == 0
            ), f"evict host leaf: host_mamba_ref_counter != 0 with {x.id=} {x.host_mamba_ref_counter=}"
            self._evict_host_leaf(x)
            num_evicted += 1
        else:
            # Internal host node
            assert (
                x.host_mamba_ref_counter == 0
            ), f"evict host mamba internal: host_mamba_ref_counter != 0 with {x.id=} {x.host_mamba_ref_counter=}"
            self.mamba_host_lru_list.remove_node(x)
            self.mamba_pool_host.free(x.mamba_host_value)
            x.mamba_host_value = None

설명:

  • 호스트 Mamba 캐시를 evict할 때, 해당 노드가 evictable leaf인지 internal node인지에 따라 다른 로직을 수행합니다. Leaf 노드의 경우 _evict_host_leaf를 호출하기 전에 host_ref_counter와 host_mamba_ref_counter가 모두 0임을 assert합니다. Internal node의 경우 Mamba 캐시만 해제하고 LRU에서 제거합니다. 이는 Mamba 캐시의 evict 로직을 더 세밀하게 제어합니다.

write_backup_storage 및 prefetch_from_storage 메소드:

백업 저장 및 사전 로딩 시 Mamba 캐시 보호 여부를 명시적으로 처리하도록 변경되었습니다.

Before (write_backup_storage):

        self.ongoing_backup[operation_id] = node
        self._protect_host_node(node)

After (write_backup_storage):

        mamba_host_protected = extra_pools is not None
        self.ongoing_backup[operation_id] = (node, mamba_host_protected)
        self._protect_host_node(node, protect_mamba=mamba_host_protected)

Before (prefetch_from_storage):

        self._protect_host_node(last_host_node)
        # ... (KV allocation) ...
        # ... (Mamba allocation) ...
        self._release_host_node(last_host_node)

After (prefetch_from_storage):

        self._protect_host_node(last_host_node, protect_mamba=False)
        # ... (KV allocation) ...
        # ... (Mamba allocation) ...
        # mamba is also being loaded, protect host mamba as well
        last_host_node.protect_host_mamba()
        if self.mamba_host_lru_list.in_list(last_host_node):
            self.mamba_host_lru_list.remove_node(last_host_node)
        self._release_host_node(last_host_node, release_mamba=False)

설명:

  • write_backup_storage에서는 백업 시 Mamba 캐시가 보호되어야 하는지 여부(mamba_host_protected)를 결정하고, 이를 ongoing_backup에 저장하며 _protect_host_node 호출 시 전달합니다.
  • prefetch_from_storage에서는 KV 캐시 할당 전에 protect_mamba=False로 _protect_host_node를 호출하여 KV 캐시만 보호합니다. 이후 Mamba 캐시가 할당되면, Mamba 캐시도 보호하고 LRU에서 제거하는 로직이 추가됩니다. 마지막으로 _release_host_node 호출 시 release_mamba=False로 하여 Mamba 캐시는 해제되지 않도록 합니다. 이는 사전 로딩 과정에서 Mamba 캐시의 상태를 정확하게 관리하기 위함입니다.

_force_release_pending_storage_ops 및 _drain_backup 메소드:

백업 작업 완료 시 Mamba 캐시 해제 여부를 처리하도록 수정되었습니다.

Before (_force_release_pending_storage_ops):

                    self._release_host_node(node)

After (_force_release_pending_storage_ops):

                    node, mamba_host_protected = entry
                    self._release_host_node(node, release_mamba=mamba_host_protected)

설명:

  • ongoing_backup에 저장된 항목이 (node, mamba_host_protected) 튜플 형태로 변경됨에 따라, 백업 작업 완료 시 _release_host_node를 호출할 때 release_mamba 인자를 전달하여 Mamba 캐시 해제 여부를 결정하도록 수정되었습니다. 이는 백업 작업이 Mamba 캐시를 사용했는지 여부에 따라 적절하게 해제되도록 합니다.

왜 이게 좋은가?

이번 PR의 변경 사항은 Mamba 모델의 호스트 캐시 관리에 있어 다음과 같은 이점을 제공합니다.

  1. 정확한 메모리 관리: Mamba 모델은 KV 캐시 외에도 자체적인 상태 벡터와 같은 추가적인 메모리를 호스트에 유지할 수 있습니다. 기존에는 이러한 Mamba 관련 메모리가 일반 KV 캐시와 동일하게 처리되어, KV 캐시만 해제되면 Mamba 관련 메모리도 함께 해제될 위험이 있었습니다. 새로운 host_mamba_ref_counter와 관련 로직은 Mamba 캐시의 참조를 명확하게 추적하고 관리함으로써, Mamba 관련 데이터가 불필요하게 해제되는 것을 방지하고 메모리 무결성을 보장합니다. 이는 특히 긴 시퀀스 처리나 복잡한 추론 과정에서 중요합니다.
  2. 효율적인 Eviction 정책: evictable_full_host_leaves 집합에 노드를 추가하는 조건에 host_mamba_ref_counter > 0이 포함됨으로써, Mamba 캐시가 사용 중인 노드는 즉시 evict 대상에서 제외됩니다. 이는 Mamba 캐시의 성능 특성을 고려한 합리적인 결정입니다. Mamba 캐시의 접근 비용이 KV 캐시보다 높을 수 있으므로, 사용 중인 Mamba 캐시를 보호하는 것은 전체적인 추론 속도에 긍정적인 영향을 미칠 수 있습니다.
  3. 성능 향상 가능성: Mamba 캐시의 참조 카운터를 세밀하게 관리함으로써, 불필요한 캐시 해제 및 재할당을 줄일 수 있습니다. 또한, Mamba 캐시가 보호되는 동안 LRU 리스트에서 제외되는 로직은 캐시 히트율을 높이고, Mamba 캐시 접근 시 발생하는 오버헤드를 줄여 전반적인 추론 속도 향상에 기여할 수 있습니다. 비록 이 PR 자체에 구체적인 성능 수치(Speed Tests and Profiling 결과)가 명시되어 있지는 않지만, 이러한 메모리 관리 개선은 LLM 서빙 시스템에서 성능 병목 현상을 완화하는 데 중요한 역할을 합니다.

일반적인 교훈:

  • 복잡한 시스템의 구성 요소별 상태 관리: LLM과 같이 복잡한 시스템에서는 다양한 종류의 캐시(KV cache, Mamba state, etc.)가 존재할 수 있습니다. 각 캐시의 특성과 의존성을 이해하고, 이를 위한 별도의 참조 카운터나 관리 메커니즘을 도입하는 것이 중요합니다. 이는 메모리 누수나 잘못된 해제를 방지하고 시스템의 안정성을 높입니다.
  • Eviction 정책의 세밀한 조정: 캐시 eviction 정책은 시스템 성능에 직접적인 영향을 미칩니다. 단순히 LRU나 LFU와 같은 일반적인 알고리즘을 적용하는 것을 넘어, 각 캐시 항목의 중요도, 접근 비용, 그리고 다른 캐시와의 상호작용을 고려하여 eviction 정책을 세밀하게 조정해야 합니다.
  • 리뷰 피드백의 중요성: 리뷰어의 피드백은 코드의 견고성과 정확성을 높이는 데 결정적인 역할을 합니다. host_ref_counter의 무결성 검증과 같이, 사소해 보일 수 있는 부분에 대한 지적은 시스템의 안정성을 크게 향상시킬 수 있습니다.

References

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.

댓글

관련 포스트

PR Analysis 의 다른글