본문으로 건너뛰기

[sglang] NixlKVManager 성능 향상: 비동기 및 멀티스레드 KV 전송 도입

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

들어가며

대규모 언어 모델(LLM)의 분산 추론 환경에서 모델의 성능은 네트워크 통신 및 데이터 전송 속도에 크게 좌우됩니다. 특히, KV 캐시와 같은 대규모 데이터를 효율적으로 전송하는 것은 추론 속도 향상의 핵심 과제입니다. 기존 sglang의 NixlKVManager는 KV 캐시 전송 요청을 동기적으로 처리하여, 특히 프리필(prefill) 단계에서 병목 현상을 야기하고 전체 처리량을 제한하는 문제가 있었습니다. 이 PR은 이러한 문제를 해결하기 위해 NixlKVManager의 KV 전송 방식을 비동기 및 멀티스레드로 개선하여, 프리필 노드에서의 성능을 획기적으로 향상시키는 것을 목표로 합니다.

코드 변경 분석

이번 PR의 핵심은 NixlKVManager의 PREFILL 모드에서 KV 캐시 전송 로직을 비동기 및 멀티스레드 방식으로 변경한 것입니다. 주요 변경 사항은 다음과 같습니다.

1. 비동기 전송을 위한 큐 및 워커 풀 도입 (sglang/srt/disaggregation/nixl/conn.py)

기존에는 add_transfer_request 함수 내에서 각 청크의 전송이 동기적으로 이루어졌습니다. 이로 인해 프리필 스케줄러가 전송 완료를 기다리며 블록되는 문제가 발생했습니다. 이를 해결하기 위해 다음과 같은 구조가 도입되었습니다.

  • FastQueue 및 ThreadPoolExecutor: 여러 개의 FastQueue 인스턴스를 생성하고, 각 큐마다 ThreadPoolExecutor를 사용하여 여러 워커 스레드가 동시에 작업을 처리할 수 있도록 합니다. 큐의 개수와 스레드 풀의 크기는 환경 변수(SGLANG_DISAGGREGATION_QUEUE_SIZE, SGLANG_DISAGGREGATION_THREAD_POOL_SIZE)를 통해 조절 가능합니다.
  • TransferKVChunk 데이터 클래스: 전송해야 할 KV 청크 정보를 담는 새로운 데이터 클래스가 정의되었습니다. 여기에는 전송할 KV 인덱스, 슬라이스 정보, 마지막 청크 여부, 청크 ID 등이 포함됩니다.
  • transfer_worker 스레드: 각 큐에서 TransferKVChunk를 가져와 실제 KV 캐시 전송(send_kvcache, send_kvcache_slice, maybe_send_extra, send_aux)을 수행하는 데몬 스레드입니다. 이 스레드는 예외 발생 시에도 계속 실행되도록 설계되었습니다.

Before: add_transfer_request 내에서 동기적으로 전송 수행

# (기존 코드 - 요약)
def add_transfer_request(...):
    # ... 전송 로직 ...
    handles = []
    for req in reqs_to_be_processed:
        # ... kv_xfer_handle = self.send_kvcache(...) ...
        handles.append(kv_xfer_handle)
    # ... handles.wait() ...

After: TransferKVChunk를 큐에 넣고 워커 스레드가 비동기적으로 처리

# (수정된 코드 - 요약)
def add_transfer_request(...):
    # ...
    kv_chunk = TransferKVChunk(...) # 전송할 청크 정보 생성
    self.transfer_queues[bootstrap_room % len(self.transfer_queues)].put(kv_chunk)
    return None # 즉시 반환

def transfer_worker(self, queue: FastQueue):
    while True:
        kv_chunk: TransferKVChunk = queue.get()
        # ... 실제 전송 로직 (send_kvcache, send_kvcache_slice 등) ...
        # ... 전송 완료 후 상태 업데이트 ...

2. add_transfer_request의 비동기화 (sglang/srt/disaggregation/nixl/conn.py)

add_transfer_request 함수는 이제 실제 전송 작업을 수행하지 않고, 생성된 TransferKVChunk 객체를 적절한 FastQueue에 넣고 즉시 None을 반환합니다. 이를 통해 호출자는 전송 완료를 기다릴 필요 없이 다음 작업을 즉시 진행할 수 있습니다.

Before:

def add_transfer_request(
    self, bootstrap_room: int, kv_indices: npt.NDArray[np.int32], ...
):
    # ... (동기 전송 로직) ...
    return handles # 전송 핸들 반환

After:

def add_transfer_request(
    self, bootstrap_room: int, kv_indices: npt.NDArray[np.int32], ...
):
    # ...
    kv_chunk = TransferKVChunk(...) # TransferKVChunk 생성
    self.transfer_queues[bootstrap_room % len(self.transfer_queues)].put(kv_chunk)
    # ...
    return None # 즉시 반환

3. NixlKVSender의 간소화 (sglang/srt/disaggregation/nixl/conn.py)

전송 핸들(xfer_handles)을 추적하고 폴링하는 로직이 제거되었습니다. 대신, kv_mgr.check_status(bootstrap_room)를 통해 요청 상태를 확인합니다. 또한, 전송 완료 후 상태를 정리하는 clear() 메서드가 추가되었습니다.

Before:

# (기존 코드 - 요약)
class NixlKVSender:
    def __init__(self):
        self.xfer_handles = {}
    
    def poll(self):
        # ... self.xfer_handles 를 폴링 ...

After:

# (수정된 코드 - 요약)
class NixlKVSender:
    # xfer_handles 제거
    
    def poll(self):
        # ... kv_mgr.check_status(bootstrap_room) 사용 ...

    def clear(self, bootstrap_room):
        # ... request_status 정리 ...

4. 스케줄러의 부트스트래핑 처리 변경 (sglang/srt/disaggregation/prefill.py)

프리필 스케줄러(prefill.py)에서 KVPoll.Bootstrapping 상태의 요청을 더 이상 완료된 것으로 간주하지 않도록 수정되었습니다. 이는 KV 전송이 완료되기 전까지는 요청이 진행 중인 것으로 처리하여, 전송 지연으로 인한 스케줄링 오류를 방지하기 위함입니다.

Before: KVPoll.Bootstrapping은 완료된 상태로 간주될 수 있었음.

After: KVPoll.Bootstrapping은 WaitingForInput, Transferring과 함께 미완료 상태로 처리됨.

# (수정된 코드 - 요약)
# In prefill.py
if req_status in [KVPoll.WaitingForInput, KVPoll.Transferring, KVPoll.Bootstrapping]:
    # ... 미완료 처리 ...

왜 이게 좋은가?

이번 변경은 NixlKVManager의 프리필 경로에서 심각한 성능 병목 현상을 제거했습니다. 이전에는 각 KV 청크 전송이 동기적으로 처리되어야 했기 때문에, 많은 디코딩 인스턴스와 청크 전송이 동시에 발생할 때 프리필 스케줄러가 전송 완료를 기다리며 블록되었습니다. 이는 전체 처리량을 제한하는 주요 원인이었습니다.

비동기 멀티스레드 전송 도입으로 다음과 같은 이점을 얻을 수 있습니다.

  • 처리량 향상: 워커 스레드가 여러 큐에서 KV 청크를 병렬로 처리하므로, 전송 작업이 더 이상 스케줄러를 블록하지 않습니다. 디코딩 인스턴스들이 여러 큐에 분산되어 처리되면서 작업 간의 오버랩이 향상됩니다.
  • 지연 시간 감소: 실제 성능 측정 결과, Qwen3-32B PD disaggregation 시나리오에서 평균 전송 시간이 162,225 μs에서 41,225 μs로 약 4배 감소했습니다. 또한, 전송 시간의 편차가 크게 줄어들어 대부분의 샘플이 34k–42k μs 대역에 분포하게 되었습니다. 이는 P95, P99와 같은 지연 시간의 꼬리 부분(tail latency)이 크게 개선되었음을 의미합니다.

이러한 개선은 분산 LLM 추론 시스템에서 데이터 전송 효율성을 높이는 것이 얼마나 중요한지를 보여주는 좋은 예시입니다. 특히, 대규모 모델과 복잡한 분산 설정에서는 이러한 최적화가 전체 시스템 성능에 지대한 영향을 미칠 수 있습니다.

리뷰 과정 및 주요 피드백

리뷰 과정에서 몇 가지 중요한 논의가 있었습니다. 특히 ShangmingCai 리뷰어는 prefill.py의 KVPoll.Bootstrapping 상태 처리에 대한 의문을 제기했습니다. pop_bootstrapped 함수가 변경되지 않은 상태에서 KVPoll.Bootstrapping 요청이 self.disagg_prefill_inflight_queue에 어떻게 들어갈 수 있는지에 대한 질문이었습니다. 이에 대해 ovidiusm은 해당 로직을 제거하여 문제를 해결했습니다. 이는 코드 변경의 정확성과 잠재적 오류를 잡아내는 리뷰 과정의 중요성을 보여줍니다.

또한, ovidiusm은 워커 스레드에서 발생하는 예외 처리에 대한 중요성을 강조했습니다. 워커 스레드가 예외로 인해 종료되면 전체 시스템이 멈출 수 있으므로, 예외를 잡아 메인 스레드로 전달하여 적절히 처리하도록 코드를 수정했습니다. 이는 _NIXL_TRANSPORT_ERRORS와 같은 특정 오류를 감지하고, 다른 예외 발생 시에도 워커 스레드가 계속 실행되도록 보장하는 안정성 확보 방안입니다.

결론

이번 PR은 NixlKVManager의 KV 전송 방식을 비동기 멀티스레드로 전환함으로써 프리필 노드의 성능을 크게 향상시켰습니다. 평균 전송 시간을 4배 가까이 줄이고 지연 시간의 편차를 감소시킨 것은 분산 LLM 추론 시스템의 효율성을 높이는 데 중요한 기여를 했습니다. 이는 복잡한 분산 시스템에서 병목 현상을 식별하고 비동기 및 병렬 처리 패턴을 적용하는 것이 얼마나 효과적인지를 잘 보여주는 사례입니다.

References

  • NixlKVManager - NixlKVManager의 구현체입니다.
  • FastQueue - 비동기 처리를 위한 큐 구현입니다.
  • ThreadPoolExecutor - Python 표준 라이브러리의 스레드 풀 실행기입니다.
  • KVPoll Status - KV 캐시 요청의 상태를 나타내는 열거형입니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글