[sglang] SGLang NIXL HiCache 리팩토링 및 O_DIRECT 지원 추가: 성능 향상과 안정성 강화

2026년 6월 1일수정: 2026년 6월 1일

PR 링크: sgl-project/sglang#25173 상태: Merged | 변경: +1561 / -835

들어가며

대규모 언어 모델(LLM)의 추론 성능을 극대화하기 위해 SGLang은 KV 캐시 관리에 있어 다양한 최적화 기법을 도입하고 있습니다. 이번 PR은 SGLang의 NIXL HiCache 커넥터의 핵심 로직을 재구성하고, 특히 파일 기반 백엔드에 대한 O_DIRECT I/O 지원을 추가하여 성능과 안정성을 동시에 향상시키는 것을 목표로 합니다.

기존 NIXL 커넥터는 스레드 동기화 문제와 객체 저장소 백엔드에서의 버그를 포함하고 있었습니다. 또한, O_DIRECT 옵션은 OS 페이지 캐시를 우회하여 NVMe와 같은 스토리지 장치에서 상당한 I/O 대역폭 향상을 기대할 수 있음에도 불구하고 지원되지 않았습니다. 이 PR은 이러한 문제들을 해결하고, 더 나아가 메모리 할당 방식 개선을 통해 LLM 추론 시 KV 캐시 로딩 성능을 한 단계 끌어올립니다.

코드 분석

1. 환경 변수 및 설정 개선 (`docs_new/docs/references/environment_variables.mdx`, `python/sglang/srt/environ.py`)

새로운 환경 변수 SGLANG_HICACHE_NIXL_USE_DIRECT_IO와 SGLANG_HUGEPAGE_SIZE가 추가되었습니다. SGLANG_HICACHE_NIXL_USE_DIRECT_IO는 NIXL 백엔드에서 O_DIRECT 사용 여부를 제어하며, 기본값은 true입니다. SGLANG_HUGEPAGE_SIZE는 호스트 KV 캐시 할당 시 Huge Pages 사용을 지원하여 메모리 관리 효율성을 높입니다.

Before:

--- a/docs_new/docs/references/environment_variables.mdx
+++ b/docs_new/docs/references/environment_variables.mdx
@@ -832,6 +832,16 @@
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Decode-side incremental KV cache offload stride. Rounded down to a multiple of <code>--page-size</code> (min is <code>--page-size</code>). If unset/invalid/&lt;=0, it falls back to <code>--page-size</code>.</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Not set (uses <code>--page-size</code>)</td>
     </tr>
+    <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_HICACHE_NIXL_USE_DIRECT_IO</code></td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Enable <code>O_DIRECT</code> for any file-based NIXL backend (POSIX, GDS, GDS_MT, 3FS) when opening cache files (bypasses the OS page cache, reducing memory pressure and improving throughput on NVMe). Can also be disabled via <code>{'{"use_direct_io": false}'}</code> in <code>--hicache-storage-backend-extra-config</code>. Falls back to buffered I/O with a warning when <code>O_DIRECT</code> is unavailable on the current OS.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>true</code></td>
+    </tr>
+    <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_HUGEPAGE_SIZE</code></td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Use huge pages for host KV cache allocations (HiCache / disaggregation offload). Valid values: <code>2MB</code> (2 MiB pages via <code>MAP_HUGE_2MB</code>) or <code>1GB</code> (1 GiB pages via <code>MAP_HUGE_1GB</code>). Requires huge pages to be pre-allocated on the host OS (<code>/proc/sys/vm/nr_hugepages</code> or <code>/sys/kernel/mm/hugepages</code>). If the allocation fails, the allocator logs a warning and falls back to regular page-size mmap automatically.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Not set (uses OS default page size)</td>
+    </tr>
   </tbody>
 </table>

After:

--- a/python/sglang/srt/environ.py
+++ b/python/sglang/srt/environ.py
@@ -333,6 +333,11 @@ class Envs:
     SGLANG_HICACHE_DECODE_OFFLOAD_STRIDE = EnvInt(None)
     SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR = EnvStr(None)
     SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR = EnvStr(None)
+    # Enable O_DIRECT when opening NIXL POSIX backend files (bypasses OS page cache).
+    # Disable with SGLANG_HICACHE_NIXL_USE_DIRECT_IO=0 or via the
+    # "use_direct_io": false key in --hicache-storage-backend-extra-config.
+    SGLANG_HICACHE_NIXL_USE_DIRECT_IO = EnvBool(True)
+    SGLANG_HUGEPAGE_SIZE = EnvStr("")
     # Staging buffer for heterogeneous TP KV transfer
     SGLANG_DISAGG_STAGING_BUFFER = EnvBool(False)
     SGLANG_DISAGG_STAGING_BUFFER_SIZE_MB = EnvInt(64)

2. 스토리지 배치 IO 개선 (`python/sglang/srt/managers/cache_controller.py`, `python/sglang/srt/mem_cache/hicache_storage.py`)

기존에는 self.storage_batch_size가 CacheController 내부에 정의되어 있었으나, 이제는 hicache_storage.py에서 STORAGE_BATCH_SIZE = 128로 상수로 정의되어 관리됩니다. 이는 스토리지 IO 작업의 일관성을 높이고 코드의 가독성을 개선합니다.

Before:

--- a/python/sglang/srt/managers/cache_controller.py
+++ b/python/sglang/srt/managers/cache_controller.py
@@ -507,8 +508,6 @@ def attach_storage_backend(
                 0, int(0.8 * (self.mem_pool_host.size - self.mem_pool_device.size))
             )
-            # granularity of batch storage IO operations, in number of pages
-            self.storage_batch_size = 128
             # tracking the number of tokens locked in prefetching, updated by the main scheduler thread
             self.prefetch_tokens_occupied = 0

After:

--- a/python/sglang/srt/mem_cache/hicache_storage.py
+++ b/python/sglang/srt/mem_cache/hicache_storage.py
@@ -16,6 +16,9 @@
 
 logger = logging.getLogger(__name__)
 
+# Max pages per batched storage IO call.
+STORAGE_BATCH_SIZE = 128
+
 
 @dataclass
 class HiCacheStorageConfig:

또한, _page_transfer 및 _storage_hit_query, _page_backup 함수 내에서 배치 처리 시 self.storage_batch_size 대신 STORAGE_BATCH_SIZE 상수를 사용하도록 변경되었습니다.

Before:

--- a/python/sglang/srt/managers/cache_controller.py
+++ b/python/sglang/srt/managers/cache_controller.py
@@ -915,8 +913,8 @@ def _page_transfer(self, operation):
     def _page_transfer(self, operation):
         # Transfer batch by batch
         prefix_keys = operation.prefix_keys
-        for i in range(0, len(operation.hash_value), self.storage_batch_size):
-            batch_hashes = operation.hash_value[i : i + self.storage_batch_size]
+        for i in range(0, len(operation.hash_value), STORAGE_BATCH_SIZE):
+            batch_hashes = operation.hash_value[i : i + STORAGE_BATCH_SIZE]
             batch_host_indices = operation.host_indices[
                 i * self.page_size : (i + len(batch_hashes)) * self.page_size
             ]
@@ -978,11 +976,9 @@ def _storage_hit_query(self, operation) -> tuple[list[str], int]:
         hash_value = []
 
         for start in range(
-            0, len(tokens_to_fetch), self.page_size * self.storage_batch_size
+            0, len(tokens_to_fetch), self.page_size * STORAGE_BATCH_SIZE
         ):
-            end = min(
-                start + self.page_size * self.storage_batch_size, len(tokens_to_fetch)
-            )
+            end = min(start + self.page_size * STORAGE_BATCH_SIZE, len(tokens_to_fetch))
             batch_tokens = tokens_to_fetch[start:end]
             batch_hashes = []
             for i in range(0, len(batch_tokens), self.page_size):
@@ -1120,8 +1116,8 @@ def _draft_page_get(self, hash_values, host_indices) -> None:
     def _page_backup(self, operation):
         # Backup batch by batch
         prefix_keys = operation.prefix_keys
-        for i in range(0, len(operation.hash_value), self.storage_batch_size):
-            batch_hashes = operation.hash_value[i : i + self.storage_batch_size]
+        for i in range(0, len(operation.hash_value), STORAGE_BATCH_SIZE):
+            batch_hashes = operation.hash_value[i : i + STORAGE_BATCH_SIZE]
             batch_host_indices = operation.host_indices[
                 i * self.page_size : (i + len(batch_hashes)) * self.page_size
             ]

3. `O_DIRECT` 지원 및 페이지 정렬 (`python/sglang/srt/mem_cache/memory_pool_host.py`)

O_DIRECT 사용 시 파일 시스템의 페이지 캐시를 우회하므로, I/O 요청 시 전달되는 메모리 주소가 OS 페이지 크기(기본 4KB)의 배수로 정렬되어야 합니다. 이를 위해 다음과 같은 변경이 이루어졌습니다.

HostTensorAllocator 변경: torch.empty 대신 alloc_mmap 함수를 사용하여 OS 페이지 정렬된 메모리를 할당합니다. 이는 kv_buffer.data_ptr()가 항상 페이지 정렬되도록 보장합니다.
HostKVCache.is_stride_page_aligned() 추가: 각 KV 캐시 풀 레이아웃(MHA, MLA 등)에 대해 스트라이드(stride)가 페이지 크기의 배수인지 확인하는 메서드가 추가되었습니다. O_DIRECT 사용 시 이 메서드가 True를 반환해야 제로 카피(zero-copy) I/O가 가능하며, 그렇지 않으면 안전을 위해 복사 모드로 폴백(fallback)합니다.

Before:

--- a/python/sglang/srt/mem_cache/memory_pool_host.py
+++ b/python/sglang/srt/mem_cache/memory_pool_host.py
@@ -37,6 +37,7 @@
     MHATokenToKVPool,
     MLATokenToKVPool,
 )
+from sglang.srt.mem_cache.mmap_allocator import alloc_mmap
 from sglang.srt.utils import is_cuda, is_hip, is_mps, is_npu, is_xpu
 
 _is_cuda = is_cuda()
@@ -78,18 +79,19 @@ def wrapper(self, *args, **kwargs):
     return wrapper
 
 
-class HostTensorAllocator(abc.ABC):
+class HostTensorAllocator:
     def __init__(self):
         """Initialize the HostTensorAllocator."""
         self.dtype = None
         self.dims = None
 
     def allocate(self, dims: tuple, dtype: torch.dtype, device: str) -> torch.Tensor:
-        """Allocate a tensor of given dims and dtype on the memory."""
+        assert (
+            device == "cpu"
+        ), f"HostTensorAllocator only supports CPU allocations; got device={device!r}"
         self.dtype = dtype
         self.dims = dims
-        tensor = torch.empty(dims, dtype=dtype, device=device)
-        return tensor
+        return alloc_mmap(dims, dtype)
 
 
 class HiSparseHostPoolMixin:
@@ -187,11 +189,16 @@ def alloc_with_host_register(
     """
     buffer = allocator.allocate(dims, dtype=dtype, device=device)
     if pin_memory:
-        ret = torch.cuda.cudart().cudaHostRegister(
-            buffer.data_ptr(), buffer.numel() * buffer.element_size(), 0
-        )
-        if ret != 0:
-            raise RuntimeError(f"cudaHostRegister failed with error code {ret}")
+        cudart = torch.cuda.cudart()
+        n_bytes = buffer.numel() * buffer.element_size()
+        rc = cudart.cudaHostRegister(buffer.data_ptr(), n_bytes, 0)
+        if int(rc) != 0:
+            raise RuntimeError(
+                f"cudaHostRegister failed (rc={int(rc)}, "
+                f"{cudart.cudaGetErrorString(rc)}) for ptr={buffer.data_ptr():#x} "
+                f"size={n_bytes}; host buffer is not pinned and device transfers "
+                f"may silently return stale data."
+            )
     return buffer
 
 
@@ -324,6 +331,19 @@ def set_from_flat_data_page(self, index: int, data_page: torch.Tensor) -> None:
         """
         raise NotImplementedError()
 
+    def is_stride_page_aligned(self, page_size_bytes: int = 4096) -> bool:
+        """Return True if per-page strides are multiples of *page_size_bytes*.
+
+        Subclasses should override this with a layout-specific stride formula.
+        This base implementation logs a warning and returns False (safe default).
+        """
+        logger.warning(
+            "%s does not implement is_stride_page_aligned(); assuming not aligned. "
+            "O_DIRECT with a file-based NIXL backend will fall back to copy mode for this pool.",
+            type(self).__name__,
+        )
+        return False
+
 
 @synchronized
 def clear(self):
@@ -850,6 +870,30 @@ def get_page_buffer_meta(self, indices):
             raise ValueError(f"Unsupported layout: {self.layout}")
         return ptr_list, element_size_list
 
+    def is_stride_page_aligned(self, page_size_bytes: int = 4096) -> bool:
+        """Return True if per-page strides are multiples of *page_size_bytes*.
+
+        When O_DIRECT is used with any file-based NIXL backend, every data pointer
+        passed to the kernel must be page-aligned.  In zero-copy mode the
+        pointer for KV page ``p`` is:
+
+            base_ptr + p * page_size * layer_num * head_num * head_dim * itemsize
+
+        For this to be page-aligned (given a page-aligned ``base_ptr``) the per-page
+        stride must itself be a multiple of the OS page size.
+        """
+        if self.layout not in ("page_first", "page_first_direct", "page_head"):
+            return False
+        stride = (
+            self.page_size
+            * self.layer_num
+            * self.head_num
+            * self.head_dim
+            * self.dtype.itemsize
+        )
+        base_aligned = self.kv_buffer.data_ptr() % page_size_bytes == 0
+        return base_aligned and stride % page_size_bytes == 0
+
 
 class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
     device_pool: MLATokenToKVPool
@@ -1250,6 +1294,26 @@ def get_page_buffer_meta(self, indices):
             raise ValueError(f"Unsupported layout: {self.layout}")
         return ptr_list, element_size_list
 
+    def is_stride_page_aligned(self, page_size_bytes: int = 4096) -> bool:
+        """Return True if per-page strides are multiples of *page_size_bytes*.
+
+        When O_DIRECT is used with any file-based NIXL backend, every data pointer
+        passed to the kernel must be page-aligned.  In zero-copy mode the
+        pointer for KV page ``p`` is:
+
+            base_ptr + p * page_size * layer_num * kv_cache_dim * itemsize
+
+        For this to be page-aligned (given a page-aligned ``base_ptr``) the per-page
+        stride must itself be a multiple of the OS page size.
+        """
+        if self.layout not in ("page_first", "page_first_direct"):
+            return False
+        stride = (
+            self.pag

4. 스레드 동기화 및 객체 저장소 버그 수정

PR 설명에 따르면, 기존 NIXL 커넥터의 스레드 동기화 문제(백업 스레드와 메인 스레드 간의 batch_set_v1, batch_get_v1 사용 충돌)와 객체 저장소 백엔드 포맷팅 버그가 수정되었습니다. 이는 코드베이스 전반에 걸쳐 안정성을 높이는 중요한 변경입니다.

5. 테스트 스위트 강화

스레드 안정성 테스트(stress test) 추가
메모리 영역 등록 전송 테스트
객체 저장소 테스트(MinIO 필요)
기존 테스트 스위트 위치 수정 및 CI 연동

이러한 테스트 강화는 코드 변경의 정확성을 검증하고 향후 발생할 수 있는 회귀(regression)를 방지하는 데 기여합니다.

왜 이게 좋은가

1. `O_DIRECT`를 통한 I/O 성능 향상

O_DIRECT 옵션은 파일 I/O 시 OS 페이지 캐시를 우회합니다. 이는 특히 NVMe SSD와 같이 빠른 스토리지 장치를 사용할 때, 이중 캐싱으로 인한 오버헤드를 줄이고 I/O 대역폭을 크게 향상시킬 수 있습니다. PR에 포함된 벤치마크 결과는 O_DIRECT 활성화 시 TTFT(Time To First Token)가 유의미하게 감소하는 것을 보여줍니다. 예를 들어, URING 백엔드에서 O_DIRECT를 사용했을 때, 콜드(cold) 상태에서 URING 대비 약 1.5배, 웜(warm) 상태에서는 약 1.2배의 성능 향상을 보였습니다.

2. 페이지 정렬된 메모리 할당 (`mmap`)

O_DIRECT의 정렬 요구사항을 충족하기 위해 mmap-기반 호스트 할당자를 도입한 것은 매우 좋은 설계입니다. 이는 kv_buffer.data_ptr()가 항상 OS 페이지 크기의 배수로 정렬되도록 보장하여, O_DIRECT 사용 시 제로 카피(zero-copy) I/O를 가능하게 합니다. 만약 정렬되지 않으면 데이터 복사 오버헤드가 발생하여 성능 이점을 상쇄할 수 있습니다.

3. Huge Pages 지원

SGLANG_HUGEPAGE_SIZE 환경 변수를 통해 Huge Pages(2MB 또는 1GB)를 사용할 수 있게 된 것은 메모리 관리 측면에서 큰 이점입니다. 일반적인 4KB 페이지보다 큰 Huge Pages를 사용하면 TLB(Translation Lookaside Buffer) miss를 줄여 메모리 접근 성능을 향상시킬 수 있으며, 특히 대규모 메모리 할당 시 효과적입니다. 할당 실패 시 일반 페이지로 자동 폴백하는 기능은 안정성을 보장합니다.

4. 안정성 및 코드 품질 향상

스레드 동기화 문제 해결: 멀티스레딩 환경에서의 잠재적인 데이터 경쟁(data race) 및 동기화 오류를 수정하여 프로그램의 안정성을 높였습니다.
객체 저장소 버그 수정: 객체 저장소 백엔드 사용 시 발생하던 포맷팅 오류 등을 수정하여 다양한 스토리지 백엔드에 대한 지원을 강화했습니다.
코드 간소화 및 테스트 강화: 불필요한 코드 제거, 테스트 스위트 개선 및 CI 연동은 코드베이스의 유지보수성을 높이고 개발 생산성을 향상시킵니다.

일반적인 교훈

I/O 최적화는 스토리지 특성에 맞춰야 합니다: O_DIRECT와 같은 저수준 I/O 옵션은 OS 페이지 캐시를 우회하여 성능을 향상시킬 수 있지만, 정렬 요구사항 등 고려해야 할 사항이 많습니다. 하드웨어 특성과 OS의 동작 방식을 이해하는 것이 중요합니다.
메모리 할당 전략의 중요성: LLM과 같이 메모리 집약적인 애플리케이션에서는 메모리 할당 방식이 성능에 직접적인 영향을 미칩니다. 페이지 정렬, Huge Pages 사용 등은 성능 최적화의 핵심 요소입니다.
안정성과 성능의 균형: 새로운 기능을 추가하거나 성능을 최적화할 때는 기존의 안정성을 해치지 않도록 주의해야 합니다. 특히 멀티스레딩 환경에서는 철저한 테스트와 코드 검토가 필수적입니다.

리뷰 피드백 반영

리뷰 과정에서 몇 가지 중요한 논의가 있었습니다:

docker/Dockerfile 관련: ishandhanani 리뷰어는 Docker 이미지에 불필요한 테스트 의존성을 추가하는 것에 대해 의문을 제기했습니다. 이에 대해 lluki는 객체 저장소 백엔드 테스트를 위해 필요하다고 설명했으나, 최종적으로 해당 테스트를 CI에서 스킵 처리하고 Dockerfile에서는 제외하는 것으로 합의되었습니다. 이는 불필요한 이미지 크기 증가를 방지하고 CI 구성을 간결하게 유지하는 데 도움이 됩니다.
주석 및 예외 처리: ishandhanani는 python/sglang/srt/mem_cache/storage/nixl/hicache_nixl.py 파일 내의 불필요한 주석 제거, O_DIRECT 정렬 요구사항에 대한 명확한 설명 추가, 그리고 폐기된 메서드에 대한 Exception 발생 처리를 요청했습니다. lluki는 이러한 피드백을 반영하여 코드의 명확성과 견고성을 높였습니다. 특히 O_DIRECT의 4KB 정렬 요구사항에 대한 설명은 사용자가 이 옵션을 올바르게 이해하고 활용하는 데 도움을 줄 것입니다.
폴링(Polling) 방식: lluki는 NIXL 백엔드와의 인터페이스에서 폴링 방식이 최적은 아니지만, 현재로서는 가장 현실적인 방법이라고 언급했습니다. 향후 더 나은 알림 메커니즘이 NIXL에 추가되거나, 성능 개선이 미흡할 경우 비동기 인터페이스를 고려할 수 있음을 시사했습니다. 이는 향후 SGLang의 HiCache 백엔드 개선 방향에 대한 중요한 단서를 제공합니다.

결론

이번 PR은 SGLang의 NIXL HiCache 커넥터를 체계적으로 개선하여 성능과 안정성을 크게 향상시켰습니다. O_DIRECT 지원, 페이지 정렬된 메모리 할당, Huge Pages 지원 등은 LLM 추론 시 KV 캐시 로딩 성능을 극대화하는 데 기여할 것입니다. 또한, 스레드 동기화 및 객체 저장소 버그 수정, 테스트 강화는 코드베이스의 전반적인 품질을 높였습니다. 이 PR은 I/O 최적화, 메모리 관리, 그리고 안정성 확보의 중요성을 잘 보여주는 사례입니다.

참고 자료

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

PR Analysis 의 다른글

이전글 [vllm] vLLM의 FP8 Scaled MM 최적화: Padding 제거를 통한 20% 성능 향상
현재글 : [sglang] SGLang NIXL HiCache 리팩토링 및 O_DIRECT 지원 추가: 성능 향상과 안정성 강화
다음글 [onnxruntime] ONNX Runtime CUDA Graph: 진정한 비동기 추론을 위한 동기화 지점 제거

[sglang] SGLang NIXL HiCache 리팩토링 및 O_DIRECT 지원 추가: 성능 향상과 안정성 강화

들어가며

코드 분석

1. 환경 변수 및 설정 개선 (`docs_new/docs/references/environment_variables.mdx`, `python/sglang/srt/environ.py`)

2. 스토리지 배치 IO 개선 (`python/sglang/srt/managers/cache_controller.py`, `python/sglang/srt/mem_cache/hicache_storage.py`)

3. `O_DIRECT` 지원 및 페이지 정렬 (`python/sglang/srt/mem_cache/memory_pool_host.py`)

4. 스레드 동기화 및 객체 저장소 버그 수정

5. 테스트 스위트 강화

왜 이게 좋은가

1. `O_DIRECT`를 통한 I/O 성능 향상

2. 페이지 정렬된 메모리 할당 (`mmap`)

3. Huge Pages 지원

4. 안정성 및 코드 품질 향상

일반적인 교훈

리뷰 피드백 반영

결론

참고 자료

댓글

관련 포스트

PR Analysis 의 다른글