[sglang] Ascend NPU 환경에서 HiCache L2 I/O 성능 최적화: Memfabric과 AscendC 활용

PR 링크: sgl-project/sglang#38826 상태: Merged | 변경: +589 / -42

들어가며

Ascend NPU 환경에서 SGLang의 HiCache(Host DRAM 기반 L2 캐시)는 호스트 메모리와 디바이스 메모리 간의 빈번한 데이터 전송을 필요로 합니다. 기존 구현체는 torch.pin_memory를 사용한 표준 호스트 메모리 할당과 aclrtStreamSynchronize를 동반하는 동기식 H2D(Host-to-Device) 전송 방식을 사용했습니다. 이 방식은 두 가지 치명적인 병목을 유발합니다.

1. AIV(Ascend Intelligent Vector) 비호환성: 기본 호스트 버퍼는 SDMA 엔진만 접근 가능하며, 고성능 offload.kv_exchange_copy 커널이 요구하는 Memfabric 매핑된 가상 주소 공간을 지원하지 않습니다.
2. 스트림 동기화 병목: 페이지 단위의 .to(device) 호출은 현재 스트림의 모든 작업을 강제로 동기화(drain)하여 파이프라인을 직렬화하고 연산 유닛을 유휴 상태로 만듭니다.

본 PR은 Memfabric 기반의 메모리 할당과 AscendC 기반의 비동기 sparse-copy IO 경로를 도입하여 이러한 문제를 해결합니다.

코드 분석

1. pool_host/mla.py: Memfabric 기반 할당 및 AscendC 커널 통합

MLATokenToKVPoolHost 클래스에 Memfabric 할당 경로를 추가하고, _transfer_ascendc_sparse_copy 메서드를 통해 퓨전 커널을 호출합니다.

# After: Memfabric을 통한 비동기 sparse copy 경로 추가
def _transfer_ascendc_sparse_copy(self, ...):
    # Fused kernel offload.kv_exchange_copy 호출
    # 호스트 포인터를 DVA(Device Virtual Address)로 변환하여 스트림 동기화 없이 전송
    offload.kv_exchange_copy(
        ...,
        stream=torch.npu.current_stream().npu_stream
    )

2. managers/cache_controller.py: 스트림 동기화 제거

기존의 .cpu() 동기화 호출 대신 to_device_no_sync를 사용하여 토큰 인덱스를 디바이스에 유지한 채 전송합니다.

# Before: .cpu() 호출로 인한 동기화 발생
return host_indices, device_indices.cpu()

# After: 비동기 전송으로 파이프라인 유지
if host_indices.device != self.device:
    host_indices = to_device_no_sync(host_indices, self.device)
return host_indices, device_indices

왜 이게 좋은가

이번 최적화의 핵심은 '데이터 전송의 비동기화'와 '커널 퓨전'입니다.

• 성능 향상: 프로파일링 결과, 기존의 aclrtStreamSynchronize로 인해 발생하던 긴 유휴 구간(Gap)이 제거되었습니다. 전송 커널이 디바이스와 호스트 메모리 사이에서 직접 데이터를 복사하면서, CPU의 개입을 최소화하고 GPU/NPU 연산 파이프라인의 흐름을 끊지 않습니다.
• 교훈: 하드웨어 가속기(NPU/GPU) 환경에서 호스트-디바이스 간 데이터 이동은 단순히 대역폭의 문제가 아니라, '동기화 지점(Sync Point)을 어디서 제거할 것인가'가 성능의 핵심입니다. Memfabric과 같은 전용 메모리 매핑 라이브러리를 활용하면 커널 수준에서 직접 메모리에 접근하여 불필요한 복사와 동기화를 방지할 수 있습니다.

리뷰어 피드백 반영

리뷰 과정에서 코드의 가독성과 유지보수를 위해 불필요한 리팩토링(변수명 변경 등)을 제거하고, ascendc_io_enabled와 같은 유틸리티 함수를 lazy import하여 모듈 의존성을 최소화하는 방향으로 수정되었습니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.Tensor.pin_memory.html
• https://docs.sglang.ai/developer_guide/contribution_guide.html

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