본문으로 건너뛰기

[SGLang] Sampler: logits에서 토큰까지의 샘플링 파이프라인

들어가며

LLM 추론의 마지막 단계는 모델이 출력한 logits 벡터에서 다음 토큰을 선택하는 샘플링이다. SGLang의 Sampler는 이 과정을 배치 단위로 최적화하며, greedy/random/top-k/top-p/min-p 등 다양한 전략을 단일 파이프라인에서 처리한다.

이 글에서는 python/sglang/srt/layers/sampler.py를 중심으로 Sampler의 구조를 분석한다.

전체 파이프라인 구조도

logits (B, V)
    │
    ▼
┌─────────────────────────────────┐
│  _preprocess_logits             │
│  ├─ Custom Logit Processor 적용 │
│  └─ NaN 감지 및 대체            │
└─────────────┬───────────────────┘
              │
    ┌─────────┴──────────┐
    │ is_all_greedy?     │
    ├─ Yes ──────────────┤──── torch.argmax → token_ids
    └─ No ───────────────┘
              │
    ┌─────────┴──────────┐
    │ simple_case?       │  (no top-k, top-p, min-p)
    ├─ Yes ──────────────┤──── softmax → multinomial
    └─ No ───────────────┘
              │
    ┌─────────┴──────────┐
    │ Backend Selection   │
    │ flashinfer/pytorch  │
    │ /ascend             │
    └─────────┬──────────┘
              │
              ▼
    batch_next_token_ids

핵심 코드 분석

Sampler 초기화

Sampler는 nn.Module을 상속하며, 서버 설정에서 NaN 감지, 결정론적 추론, RL on-policy 모드 등을 결정한다.

class Sampler(nn.Module):
    def __init__(self):
        super().__init__()
        self.use_nan_detection = get_global_server_args().enable_nan_detection
        self.tp_sync_group = get_tp_group().device_group
        self.rl_on_policy_target = get_global_server_args().rl_on_policy_target
        self.enable_deterministic = (
            get_global_server_args().enable_deterministic_inference
        )

DP Attention이 활성화된 경우 TP 동기화 그룹을 별도로 설정한다.

전처리: Custom Logit Processor와 NaN 처리

forward() 진입 직후, logits에 대한 전처리가 수행된다.

def _preprocess_logits(self, logits, sampling_info):
    if sampling_info.has_custom_logit_processor:
        apply_custom_logit_processor(logits, sampling_info)

    if self.use_nan_detection and torch.any(torch.isnan(logits)):
        logger.warning("Detected errors during sampling! NaN in the logits.")
        logits = torch.where(
            torch.isnan(logits), torch.full_like(logits, -1e5), logits
        )
    return logits

NaN이 감지되면 -1e5로 대체하여 softmax 이후 거의 0에 가까운 확률로 만든다. crash_on_warnings() 설정 시 예외를 던져 디버깅을 돕는다.

Greedy 경로 vs 샘플링 경로

배치 내 모든 요청이 greedy일 때는 단순 torch.argmax로 처리한다.

if sampling_info.is_all_greedy:
    batch_next_token_ids = torch.argmax(logits, -1)

그렇지 않으면 temperature 스케일링 후 softmax를 적용한다. 이때 메모리 절약을 위해 in-place 연산을 사용한다.

logits.div_(sampling_info.temperatures)
logits[:] = torch.softmax(logits, dim=-1)
probs = logits

백엔드별 샘플링 구현

_sample_from_probs는 flashinfer과 pytorch 두 가지 백엔드를 지원한다.

if backend == "flashinfer":
    if sampling_info.need_min_p_sampling:
        probs = top_k_renorm_prob(probs, sampling_info.top_ks)
        probs = top_p_renorm_prob(probs, sampling_info.top_ps)
        batch_next_token_ids = min_p_sampling_from_probs(
            probs, sampling_info.min_ps
        )
    else:
        batch_next_token_ids = top_k_top_p_sampling_from_probs(
            probs.contiguous(),
            sampling_info.top_ks,
            sampling_info.top_ps,
            filter_apply_order="joint",
        )

flashinfer 백엔드는 CUDA 커널로 top-k/top-p를 joint로 적용하여 성능이 우수하다. pytorch 백엔드는 순수 PyTorch 연산으로 fallback을 제공한다.

결정론적 샘플링: Gumbel Trick

결정론적 추론이 필요할 때 SGLang은 Gumbel trick과 MurmurHash를 결합한다.

@torch.compile(dynamic=True)
def multinomial_with_seed(logprobs, seed, positions):
    n, m = logprobs.shape
    seed = seed.to(torch.uint64)
    col_indices = torch.arange(m, device=logprobs.device)
    hashed = murmur_hash32(seed, positions, col_indices)

    x = hashed.to(torch.float64) / torch.iinfo(torch.uint32).max
    x.log_().clamp_(min=torch.finfo(x.dtype).min).neg_()  # -log(x)
    x.log_().neg_()  # -log(-log(x)) == gumbel noise
    x.add_(logprobs.to(torch.float64))

    return torch.argmax(x, dim=1, keepdim=True)

요청별 시드와 위치 정보로 해시를 생성하고, Gumbel 노이즈를 더해 argmax로 샘플링한다. float64를 사용하여 수치적 안정성을 보장한다.

TP 동기화

멀티 GPU 환경에서 TP rank 간 토큰 ID 불일치를 방지하기 위한 동기화 로직이 있다.

def _sync_token_ids_across_tp(self, batch_next_token_ids, sampling_info):
    if SYNC_TOKEN_IDS_ACROSS_TP or sampling_info.grammars:
        torch.distributed.all_reduce(
            batch_next_token_ids,
            op=dist.ReduceOp.MIN,
            group=self.tp_sync_group,
        )

Grammar-constrained decoding 사용 시 특히 TP rank 간 비결정성이 발생할 수 있어 항상 동기화한다.

설계 근거

설계 선택 이유
greedy 경로 분리 argmax가 softmax + multinomial보다 빠르므로 all-greedy 배치를 빠르게 처리
in-place softmax logits[:] = softmax(logits) 로 별도 텐서 할당 없이 메모리 절약
flashinfer 우선 CUDA 커널이 PyTorch 대비 top-k/top-p joint 필터링에서 성능 우위
float64 Gumbel 32비트에서는 Gumbel 노이즈의 수치적 불안정성이 결정론적 재현을 깨뜨림
TP MIN reduce all-reduce MIN으로 토큰 동기화 시 결정론적 결과 보장

관련 포스트

참고

댓글

관련 포스트

SGLang 의 다른글