FlashSampling: 빠르고 메모리 효율적인 정확한 샘플링 기법

LLM이 다음 단어를 선택할 때 수만 개의 후보 중 하나를 고르는 '샘플링' 과정은 그동안 메모리 낭비가 심한 병목 구간이었다. 이 논문은 데이터를 메모리에 썼다 읽는 과정을 생략하고 연산 과정에서 즉시 샘플링을 수행하여 추론 속도를 획기적으로 높였다. 이는 특히 실시간 응답이 중요한 챗봇이나 대규모 서비스 환경에서 비용 절감과 사용자 경험 개선에 직접적인 도움을 준다.

기존의 LLM 샘플링은 행렬 연산으로 나온 수만 개의 점수(Logits)를 일단 메모리에 다 적어두고, 다시 읽어와서 Softmax 확률을 구한 뒤 주사위를 던지는 방식이었다. 이 과정에서 거대한 데이터를 메모리에 쓰고 읽는 행위 자체가 연산보다 훨씬 오래 걸리는 병목 현상을 일으킨다. FlashSampling은 이 과정을 '주방에서 요리를 다 하고 접시에 담아 서빙하는 것'이 아니라 '요리하는 냄비 안에서 바로 가장 맛있는 조각 하나만 골라내는 것'으로 바꿨다.

핵심 원리는 Gumbel-Max 트릭이다. 딥러닝에서 확률 분포를 만드는 Softmax 연산은 모든 로짓 값을 알아야 하지만, Gumbel-Max 트릭은 각 로짓에 특수한 노이즈를 더한 뒤 가장 큰 값을 고르기만 하면 수학적으로 Softmax 결과와 완벽히 동일한 샘플을 얻을 수 있다. 이 성질 덕분에 전체 데이터를 한꺼번에 볼 필요 없이, GPU 칩 내부의 작은 메모리(SRAM)에서 타일 단위로 계산을 수행하면서 즉시 노이즈를 섞고 그 타일의 1등 후보만 남길 수 있다.

결과적으로 수만 개의 로짓 데이터 대신 각 타일의 승자 후보 수십 개만 메모리에 기록하면 된다. 마지막에 이 후보들끼리만 한 번 더 비교하면 최종 단어가 결정된다. 이는 FlashAttention이 어텐션 연산의 메모리 접근을 최적화한 것과 유사한 철학을 샘플링 영역에 적용한 것으로, 정확도를 전혀 희생하지 않으면서도 추론 효율을 극대화한다.

FlashSampling은 크게 두 단계의 퓨즈드 커널 구조로 작동한다. 첫 번째 단계인 Stage 1(Fused Kernel)에서는 GPU의 각 스레드 블록이 어휘 사전의 특정 타일을 담당한다. 입력된 히든 스테이트 $H$와 가중치 $W$의 부분 행렬을 곱해 로짓 타일을 계산하고, 여기에 온도 조절(Temperature scaling)이나 마스킹을 적용한 뒤 Gumbel 노이즈 $g_{b,i}$를 더한다. [로짓 $\ell$과 유니폼 분포에서 추출한 $u$를 입력으로 → $g = -\log(-\log u)$ 연산을 거쳐 → 노이즈가 섞인 점수 $s = \ell + g$를 얻고 → 이 중 최대값 $m$과 인덱스 $idx$를 출력한다.] 이 결과물인 타일별 승자 정보만 HBM에 기록한다.

두 번째 단계인 Stage 2(Reduction)는 각 타일에서 올라온 소수의 후보들을 모아 최종 승자를 가린다. [각 타일의 최대값 $m^{(t)}$들을 입력으로 → 다시 한번 argmax 연산을 수행해 → 전체 어휘 중 가장 높은 점수를 가진 인덱스 $i^*$를 결정하고 → 최종 샘플링된 토큰으로 반환한다.] 이 과정은 전체 로짓 텐서를 생성하고 저장하는 비용을 완전히 제거한다.

분산 환경(Tensor Parallel)에서는 각 GPU(Rank)가 자신의 로컬 어휘 파티션에서 로컬 샘플 $z_k$와 로그 질량(Log-mass) $L_k$를 계산한다. [로컬 로짓들을 입력으로 → $\log(\sum \exp(\ell_i))$ 연산을 통해 → 해당 파티션의 전체 확률 질량을 나타내는 스칼라 값 $L_k$를 얻는다.] 이후 각 GPU는 전체 로짓 대신 이 작은 요약 정보만 공유하며, 마스터 노드는 이 정보들을 바탕으로 어느 GPU의 샘플을 최종 선택할지 결정한다. 이 계층적 방식은 통신 복잡도를 $O(V)$에서 $O(\#ranks)$로 줄여준다.

커널 수준 마이크로 벤치마크 결과, H100, H200, B200, B300 등 최신 NVIDIA GPU 아키텍처 전반에서 일관된 성능 향상을 보였다. 특히 B300 GPU에서 Multinomial Sampling 대비 최대 1.84배, 기존 최적화 라이브러리인 FlashInfer의 Gumbel-Max 샘플러 대비 최대 1.39배 빠른 속도를 기록했다. 이는 샘플링 연산이 Matmul 에필로그에 통합되면서 메모리 대역폭 활용률이 극대화되었기 때문이다.

vLLM에 통합하여 수행한 엔드-투-엔드 실험에서도 실질적인 이득이 확인되었다. Qwen3-1.7B 모델을 사용한 추론 테스트에서 토큰당 생성 시간(TPOT)을 최대 19%까지 단축했다. 모델 크기가 커질수록 전체 실행 시간 중 LM-head가 차지하는 비중이 줄어들어 속도 향상 폭은 다소 감소하지만, Qwen3-8B에서도 약 3~7%의 유의미한 성능 개선을 유지했다.

정확성 검증 측면에서는 5,000개의 샘플을 대상으로 카이제곱 적합도 검정을 수행한 결과, 기존 PyTorch 구현체와 통계적으로 차이가 없음이 확인되었다. 또한 GSM8K 데이터셋을 이용한 벤치마크에서도 기존 방식(89.6%)과 FlashSampling(89.4%)이 오차 범위 내의 동일한 정확도를 보여, 최적화 과정에서 수학적 엄밀함이 유지되었음을 입증했다.

FlashSampling의 기술적 핵심은 Gumbel 분포의 '최대값 안정성(Max-stability)'과 '계층적 분해 가능성'에 있다. 논문은 임의의 파티션에 대해 글로벌 최대값이 각 타일의 로컬 최대값들의 최대값과 같다는 점을 이용하여, Matmul 연산과 동시에 타일 단위 argmax를 수행하는 IO-aware 커널 구조를 제안했다. 이는 메모리 계층 구조에서 속도가 빠른 SRAM을 최대한 활용하고 느린 HBM 접근을 최소화하는 전략이다.

이론적으로는 Group-Gumbel-Max 알고리즘을 통해 온라인 및 분산 샘플링의 정확성을 보장한다. 각 그룹의 로그 질량 $L_k = \log \sum \exp(\ell_i)$을 계산함으로써, 전체 분포를 직접 형성하지 않고도 계층적으로 정확한 샘플링이 가능함을 수학적으로 증명했다. 특히 온라인 방식에서는 새로운 그룹이 들어올 때마다 Bernoulli 선택을 통해 기존 샘플을 업데이트하는 이진 병합 규칙(Binary merge rule)을 사용하여 메모리 사용량을 $O(1)$로 유지한다.

구현 측면에서는 Triton을 사용하여 플랫폼 이식성을 확보했으며, 수치적 안정성을 위해 FP32 정밀도로 Gumbel 노이즈 생성 및 누적 연산을 수행한다. 또한 Top-k 샘플링으로의 확장이 용이하도록 설계되었으며, Top-p(Nucleus) 샘플링의 경우에도 Top-k로 후보군을 먼저 압축한 뒤 적용하는 방식으로 효율성을 확보할 수 있음을 보여주었다.

Top-p(Nucleus) 샘플링은 전체 로짓의 정렬과 누적 합계가 필요하므로 FlashSampling의 타일 기반 구조에 직접 통합하기 어렵다. 논문에서는 Top-k를 먼저 수행한 후 소수의 후보에 대해 Top-p를 적용하는 우회 방법을 제시하지만, 순수 Top-p 샘플링에 대한 직접적인 퓨전은 향후 과제로 남아 있다.