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

FlashSampling: Fast and Memory-Efficient Exact Sampling

March 16, 2026

저자: Tomas Ruiz, Zhen Qin, Yifan Zhang, Xuyang Shen, Yiran Zhong, Mengdi Wang

cs.AI

초록

카테고리 분포에서의 샘플링은 수학적으로 간단하지만, 대규모 어휘 집합을 다루는 디코딩 과정에서는 언어 모델 헤드 연산 이후 추가 메모리 트래픽과 추가 커널 실행을 자주 유발합니다. 본 논문에서는 샘플링을 LM 헤드 행렬 곱셈에 융합하고 로짓 텐서를 HBM에 물리적으로 저장하지 않는 정확한 샘플링 기법인 FlashSampling을 제안합니다. 이 방법은 단순합니다: 온칩에서 타일 단위로 로짓을 계산하고, 검블 노이즈를 추가하며, 행과 어휘 타일마다 하나의 최대값만 유지한 후 타일 간 소규모 리덕션으로 마무리합니다. 이 융합 타일 커널은 argmax가 분할 영역에서 분해 가능하기 때문에 정확합니다. 온라인 및 텐서 병렬 환경을 위한 그룹화 변형은 카테고리 분포의 계층적 인수분해를 통해 정확성을 보장합니다. H100, H200, B200, B300 GPU에서 FlashSampling은 커널 수준 디코딩 워크로드의 속도를 높이며, 종단 간 vLLM 실험에서 테스트한 모델들의 출력 토큰당 소요 시간을 최대 19%까지 단축했습니다. 이러한 결과는 근사화 없이 정확한 샘플링이 행렬 곱셈 자체에 통합될 수 있음을 보여주며, 대역폭 제한적인 후처리 단계를 가벼운 에필로그로 전환합니다. 프로젝트 페이지: https://github.com/FlashSampling/FlashSampling.

English

Sampling from a categorical distribution is mathematically simple, but in large-vocabulary decoding, it often triggers extra memory traffic and extra kernels after the LM head. We present FlashSampling, an exact sampling primitive that fuses sampling into the LM-head matmul and never materializes the logits tensor in HBM. The method is simple: compute logits tile-by-tile on chip, add Gumbel noise, keep only one maximizer per row and per vocabulary tile, and finish with a small reduction over tiles. The fused tiled kernel is exact because argmax decomposes over a partition; grouped variants for online and tensor-parallel settings are exact by hierarchical factorization of the categorical distribution. Across H100, H200, B200, and B300 GPUs, FlashSampling speeds up kernel-level decode workloads, and in end-to-end vLLM experiments, it reduces time per output token by up to 19% on the models we test. These results show that exact sampling, with no approximation, can be integrated into the matmul itself, turning a bandwidth-bound postprocessing step into a lightweight epilogue. Project Page: https://github.com/FlashSampling/FlashSampling.