[sglang] SGLang의 V4 Paged KV Cache Dequantization 성능 최적화 분석

PR 링크: sgl-project/sglang#42929 상태: Merged | 변경: +147 / -23

들어가며

LLM 추론 엔진에서 KV Cache는 메모리 대역폭을 가장 많이 소모하는 병목 지점 중 하나입니다. SGLang의 DeepSeek V4(DSV4) 레이아웃을 사용하는 paged KV 캐시 디퀀타이제이션 커널은 기존에 토큰당 하나의 프로그램을 실행하는 방식이었습니다. 이 방식은 작은 로드(small loads)를 빈번하게 발생시켜 하드웨어의 HBM 대역폭을 충분히 활용하지 못하는 문제가 있었습니다. 본 PR은 이 커널을 벡터화하여 프로그램당 4개의 토큰을 처리하도록 개선함으로써 성능을 대폭 향상시켰습니다.

코드 분석

1. 커널 호출부 최적화 (python/sglang/kernels/ops/attention/dsv4/dequant_k_cache.py)

기존에는 num_tokens만큼의 프로그램을 실행했으나, 이제는 tokens_per_prog=4를 도입하여 그리드 크기를 1/4로 줄였습니다.

Before:

_dequantize_k_cache_paged_kernel[(num_tokens,)](
    out, buf_fp8, buf_bf16, buf_uint8, page_table_1_flattened, out.stride(0), ...
)

After:

tokens_per_prog = 4
_dequantize_k_cache_paged_kernel[(triton.cdiv(num_tokens, tokens_per_prog),)](
    out, buf_fp8, buf_bf16, buf_uint8, page_table_1_flattened, out.stride(0), num_tokens, ..., TOKENS_PER_PROG=tokens_per_prog
)

2. 커널 로직 벡터화

커널 내부에서 tl.program_id를 활용해 4개의 토큰을 동시에 처리하도록 변경했습니다. tl.arange를 사용하여 데이터를 일괄 로드하고, mask를 통해 패딩된 영역이나 마지막 토큰 그룹을 안전하게 처리합니다.

After (핵심 로직):

token_ids = tl.program_id(0).to(tl.int64) * TOKENS_PER_PROG + tl.arange(0, TOKENS_PER_PROG).to(tl.int64)
tmask = token_ids < num_tokens
loc = tl.load(page_table_ptr + token_ids, mask=tmask, other=0).to(tl.int64)
# ... 이후 fp8_vals, scale_u8 등을 일괄 로드 및 연산

왜 이게 좋은가

이번 최적화의 핵심은 메모리 접근 패턴의 개선입니다. 기존의 토큰당 1개 프로그램 방식은 HBM 대역폭의 약 12% 수준밖에 활용하지 못했습니다. 4개의 토큰을 하나의 프로그램에서 처리하도록 벡터화함으로써, 더 큰 단위의 메모리 로드(Coalesced Access)가 가능해졌습니다.

• 성능 향상: B200 환경에서 262K 행을 처리할 때 기존 0.381ms에서 0.089ms로 약 4.3배의 속도 향상을 기록했습니다.
• 교훈: Triton 커널 작성 시, 단순히 연산 효율뿐만 아니라 메모리 로드/스토어의 병렬성과 단위 크기가 성능에 미치는 영향이 매우 큼을 보여줍니다. 특히 tl.constexpr과 mask를 적절히 활용하여 분기 없이 벡터화된 연산을 수행하는 것이 중요합니다.

리뷰 피드백 반영

리뷰 과정에서 environ.py에 추가된 불필요한 환경 변수 설정이 제거되었으며, 실험적인 기능(Per-ratio candidate publishing)은 리스크가 크다는 판단하에 삭제되었습니다. 또한, atol=0을 사용한 엄격한 단위 테스트를 추가하여 최적화 과정에서 발생할 수 있는 정밀도 손실을 원천 차단했습니다.

참고 자료

• https://triton-lang.org/main/index.html
• https://pytorch.org/docs/stable/generated/torch.testing.assert_close.html

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