본문으로 건너뛰기

[sglang] HMM GPU의 페이지 폴트 폭풍 피하기: prefill PLE 행을 호스트에서 모아 한 번에 전송하기

PR 링크: sgl-project/sglang#50318 상태: Merged | 변경: +326 / -10

들어가며

SGLang에서 Qwen3.8-Flash-Next는 47.7 GiB에 달하는 PLE(Per-Layer Embedding) n-gram 테이블을 사용합니다. --ple-offload-backend file 옵션을 쓰면 이 테이블을 file-backed mmap에 두고, Triton gather kernel이 GPU에서 이 메모리를 직접 읽습니다. 여기에 paged experts까지 더하면 16 GB짜리 소비자용 GPU(RTX 5070 Ti)에서도 전체 모델을 서빙할 수 있습니다.

문제는 HMM(Heterogeneous Memory Management)으로 pageable memory를 읽는 discrete GPU에서 발생합니다. GPU가 아직 매핑하지 않은 페이지에 접근할 때마다 드라이버가 page fault를 한 번에 하나씩, 약 0.5 ms씩 처리합니다. prefill은 토큰당 16개 행을 gather하고 대부분 서로 다른 페이지에 있습니다. 약 1.5k 토큰 프롬프트면 고유 행이 17~21k개이고, 처음 보는 프롬프트에서 gather kernel 하나가 9~13초를 소모했습니다. 이미 매핑된 경우 0.4초면 끝나는 작업입니다. 기존 WILLNEED 힌트는 페이지를 page cache로 끌어오긴 하지만, GPU는 여전히 각 페이지를 자기 매핑으로 하나씩 fault-in 해야 했습니다. TTFT의 대부분이 여기서 나왔습니다.

반면 GB10(DGX Spark)처럼 ATS로 호스트 page table을 직접 쓰는 GPU는 이런 개별 fault가 없으므로, 이 PR은 그 경로는 건드리지 않습니다.

코드 분석

1. qwen4_exp_ple_table.py: "이 GPU는 fault를 맞는가?" 판별

기존에는 pageable memory 접근 가능 여부만 물었습니다. 이번에 attribute 조회를 _device_attribute로 분리하고, 두 attribute를 조합한 판별 함수를 추가했습니다.

# Before
if rc == 0:
    return bool(value.value)

# After
_CUDA_DEV_ATTR_PAGEABLE_MEMORY_ACCESS_USES_HOST_PAGE_TABLES = 100

def device_faults_pageable_pages(device_index: int = 0) -> bool:
    reads = _device_attribute(_CUDA_DEV_ATTR_PAGEABLE_MEMORY_ACCESS, device_index)
    host_tables = _device_attribute(
        _CUDA_DEV_ATTR_PAGEABLE_MEMORY_ACCESS_USES_HOST_PAGE_TABLES, device_index
    )
    return reads == 1 and host_tables == 0

cudaDevAttrPageableMemoryAccess == 1이면서 cudaDevAttrPageableMemoryAccessUsesHostPageTables == 0인 경우, 즉 HMM 기반 discrete GPU만 새 경로를 탑니다. 조회 실패 시 None이 되어 False가 반환되므로 보수적으로 기존 경로를 유지합니다. 이 판별은 레이어 생성 시 한 번만 수행됩니다.

2. qwen4_exp.py: gather 분기

초기화 시 파일 경로가 있고 GPU가 fault를 맞는 경우에만 _stages_on_host를 켜고, 테이블 파일의 read-only fd를 미리 열어둡니다.

table_path = getattr(host_table, "_sglang_ple_file_path", None)
self._stages_on_host = (
    table_path is not None
    and torch.cuda.is_available()
    and device_faults_pageable_pages(torch.cuda.current_device())
)
self._table_fd = (
    os.open(table_path, os.O_RDONLY) if self._stages_on_host else None
)

gather의 핵심 변화는 kernel에 넘기는 table과 ids, 그리고 vocab 범위를 변수화한 것입니다.

# Before
_gather_ple_embedding_from_pinned_kernel[(flat_ids.numel(),)](
    self.weight.data_ptr(),
    flat_ids,
    output,
    ...
    tp_vocab_start=self.shard_indices.org_vocab_start_index,
    tp_vocab_end=self.shard_indices.org_vocab_end_index,

# After
table, ids = self.weight, flat_ids
if self._gathers_on_host(flat_ids):
    table, ids = self._stage_rows(flat_ids, vocab_start, vocab_end)
    vocab_start, vocab_end = 0, table.shape[0]
elif self._file_prefetcher is not None:
    self._file_prefetcher.enqueue(...)
_gather_ple_embedding_from_pinned_kernel[(flat_ids.numel(),)](
    table.data_ptr(), ids, output, ...,
    tp_vocab_start=vocab_start, tp_vocab_end=vocab_end, ...)

주목할 점은 같은 Triton kernel을 그대로 재사용한다는 것입니다. staging된 행 버퍼를 "작은 테이블"로 보고 id만 그 안으로 remap하면, FP8/bf16 변환 로직이 동일하므로 출력 바이트가 정확히 같습니다. 새 kernel을 작성하지 않아 정확성 검증 부담이 크게 줄어듭니다.

호스트 경로 진입 조건은 엄격합니다.

return (
    self._stages_on_host
    and flat_ids.is_cuda
    and flat_ids.numel() >= PLE_FILE_PREFETCH_MIN_ROWS  # 2048
    and not torch.cuda.is_current_stream_capturing()
)

decode 크기 gather와 CUDA graph capture 중에는 절대 진입하지 않습니다. _stage_rows는 .cpu()로 동기화하므로 graph capture와 양립할 수 없고, decode에서 동기화를 넣으면 손해이기 때문입니다.

3. _stage_rows: 고유 행만 pinned memory로

ids = flat_ids.cpu()
in_range = (ids >= vocab_start) & (ids < vocab_end)
rows, inverse = torch.unique(ids[in_range] - vocab_start, return_inverse=True)
if rows.numel() == 0:
    rows = torch.zeros(1, dtype=torch.int64)
table = self.weight.data.view(torch.uint8)
advise_rows_willneed(self._table_fd, rows, table.shape[1])
staged = torch.empty((rows.numel(), table.shape[1]), dtype=torch.uint8, pin_memory=True)
torch.index_select(table, 0, rows, out=staged)
local = torch.full(ids.shape, -1, dtype=torch.int64)
local[in_range] = inverse
  • torch.unique(return_inverse=True)로 중복을 제거하고 원래 위치 → staged 인덱스 매핑을 한 번에 얻습니다.
  • TP shard 밖의 id는 -1로 남겨 kernel이 여전히 0 행을 출력하게 합니다.
  • shard 안에 id가 하나도 없을 때는 kernel이 "범위 밖 id도 row 0을 읽은 뒤 0으로 만든다"는 동작 때문에 더미 행 하나를 둡니다. 세심한 edge case 처리입니다.
  • view(torch.uint8)로 dtype과 무관하게 바이트 단위로 복사하고, GPU로 보낸 뒤 원래 dtype으로 view합니다.

4. advise_rows_willneed: 단일 스레드 복사를 병렬 I/O로

서버는 torch를 1 스레드로 돌리고 매핑은 MADV_RANDOM이라, 그냥 index_select하면 cold 행마다 4 KiB 읽기가 직렬로 일어납니다. 그래서 복사 전에 posix_fadvise(WILLNEED)로 읽기를 모두 큐에 넣습니다.

first = rows.to(torch.int64) * row_bytes
pages = torch.cat(
    [first >> _PAGE_SHIFT, (first + row_bytes - 1) >> _PAGE_SHIFT]
).unique(sorted=True)
run_starts = torch.ones_like(pages, dtype=torch.bool)
run_starts[1:] = pages[1:] != pages[:-1] + 1
starts = torch.nonzero(run_starts).flatten()
lengths = torch.diff(starts, append=starts.new_tensor([pages.numel()]))
for page, count in zip(pages[starts].tolist(), lengths.tolist()):
    os.posix_fadvise(fd, page << _PAGE_SHIFT, count << _PAGE_SHIFT, os.POSIX_FADV_WILLNEED)

두 페이지에 걸치는 행은 시작/끝 페이지를 모두 포함하고, 인접 페이지는 run으로 병합해 syscall 수를 줄입니다. 이 run 계산이 Python 루프가 아닌 tensor 연산으로 벡터화된 점도 좋습니다. 실패하면 한 번만 경고를 남기고 중단하는 advisory 동작이라 정확성에는 영향이 없습니다. 흥미롭게도 초기 버전은 mapping에 madvise를 걸었지만 최종 버전은 별도 fd에 posix_fadvise를 써서 프롬프트당 ~10 ms를 더 줄였습니다.

왜 이게 좋은가

RTX 5070 Ti(16 GB), Qwen3.8-Flash-Next-FP8 기준:

main this PR
TTFT, ~1.5k 토큰, 첫 접근 17.10 s 7.42 s
TTFT, 반복 프롬프트 7.82 s 7.18 s
decode (tok/s) 14.93 14.79 (노이즈 범위)

gather만 따로 측정하면 차이가 더 극적입니다.

첫 접근 반복
kernel 직접 읽기 9.3–13.2 s 0.38–0.47 s
호스트 복사, 힌트 없음 0.83–1.22 s 0.002 s
호스트 복사 + WILLNEED 0.039–0.060 s 0.010–0.012 s

첫 접근 TTFT가 57% 감소했고 gather 자체는 200배 이상 빨라졌습니다. GSM8K 정확도도 0.98로 유지됐습니다.

일반적 교훈

  1. "GPU가 호스트 메모리를 직접 읽을 수 있다"는 것이 "빠르다"는 뜻은 아니다. HMM의 fault 경로는 직렬이고 페이지당 비용이 큽니다. 접근이 넓게 흩어져 있다면 CPU에서 모아 한 번의 bulk DMA로 보내는 편이 훨씬 낫습니다.
  2. 하드웨어 특성으로 분기하라. ATS(GB10)와 HMM은 같은 API 뒤에 숨어 있지만 비용 모델이 전혀 다릅니다. device attribute로 정확히 구분해 측정한 플랫폼에만 최적화를 적용했습니다.
  3. kernel을 재사용하고 입력만 바꿔라. id remap으로 기존 kernel을 그대로 쓰면 byte-identical 출력이 자연스럽게 보장됩니다.
  4. 단일 스레드 I/O라면 readahead 힌트로 병렬성을 확보하라. posix_fadvise(WILLNEED) 하나가 0.81.2초를 4060 ms로 줄였습니다.
  5. hot path(decode, CUDA graph)는 건드리지 않는다. 크기 임계값과 capture 체크로 새 코드가 decode에 절대 개입하지 않게 했습니다.

또한 negative control 테스트(device 체크를 무시하거나 WILLNEED를 제거하면 정확히 해당 테스트만 실패)를 함께 제시해 테스트 자체의 유효성까지 입증한 점은 성능 PR의 좋은 모범입니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글