How Deep Learning Computes

nano-vLLM Sampling Kernel Candidate

마지막 수정:

inferencenano-vllmsamplingkernel

Sampling은 attention보다 작아 보이지만 serving에서는 자주 호출된다.

Decode 단계에서는 매 token마다 logits processing과 token selection이 일어난다.

logits
  -> temperature
  -> top-k / top-p
  -> sample or argmax
  -> next token

처음에는 greedy argmax로 충분하다. 하지만 top-k/top-p, repetition penalty, structured decoding이 붙으면 sampling 경로도 무시하기 어려워질 수 있다.

이 카드는 sampling이 진짜 병목인지 profiler로 확인한 뒤 kernel 후보로 올릴지 결정하는 카드다.

확인

  • Sampling은 prefill보다 decode에서 더 자주 호출되는 이유는 무엇인가?
  • Greedy sampling과 top-p sampling은 계산 복잡도가 어떻게 다른가?
  • Sampling kernel 최적화는 언제 attention 최적화보다 우선순위가 낮은가?
먼저 보면 좋은 카드nano-vLLM Sampling과 Streaminglogits에서 next token을 고르고 request output에 반영하는 sampling/streaming 경로를 구현한다.먼저 보면 좋은 카드nano-vLLM Kernel Integration Pointsnano-vLLM serving loop에서 RMSNorm, attention, sampling 같은 kernel 교체 지점을 찾는다.같은 경로의 이전 카드nano-vLLM Attention Kernel Swapmanual/SDPA/optimized attention을 nano-vLLM prefill/decode loop에 연결하고 전후 성능을 비교한다.같은 경로의 다음 카드nano-vLLM Prefill/Decode Benchmark Reportkernel swap 전후의 TTFT, TPOT, throughput, memory를 prefill/decode workload별로 정리한다.