[sglang] EAGLE Speculative Decoding의 숨은 비용 제거: Vocab Shard Verify, Post-Verify CUDA Graph, 한 번의 Launch로 Draft 메타데이터 갱신

PR 링크: sgl-project/sglang#47200 상태: Merged | 변경: +4808 / -363

들어가며

Speculative decoding(EAGLE, MTP 등)은 이론상 "draft 모델이 여러 토큰을 제안하고 target 모델이 한 번에 검증"하는 단순한 구조입니다. 하지만 실제 서빙 런타임에서 한 라운드를 들여다보면 GPU가 큰 matmul을 하는 시간보다 그 사이의 잡일에 쓰는 시간이 상당합니다.

• verify 단계에서 tensor-parallel rank들의 vocabulary logits를 전부 all-gather한 뒤 argmax
• sampling → acceptance → output 준비로 이어지는 eager tail의 여러 개의 작은 kernel launch
• draft step마다 attention metadata(cache_seqlens, cu_seqlens_k)를 step 수만큼 반복 갱신
• staging metadata를 host로 동기 readback

배치가 작을수록(TP8 B1 등) 이런 launch overhead와 동기화가 지배적인 비용이 됩니다. 이 PR은 SGLANG_ENABLE_NV_MEGA_RUNTIME_OPTIMIZATION 마스터 스위치 뒤에 모델 독립적인 fast path들을 추가하고, 스위치가 꺼져 있으면 main과 동일하게 동작하도록 설계되어 있습니다. 최종 리비전은 run-to-run noise 이내였던 4개 변경을 제거하고 효과가 측정된 커밋만 남겼다는 점도 눈여겨볼 만합니다.

코드 분석

1. trtllm_mha_graph_metadata.py — draft step 메타데이터를 한 번의 launch로

기존에는 draft decode step s마다 update_trtllm_mha_graph_metadata를 seqlen_offset = first_offset + s로 호출해야 했습니다. 즉 step이 N개면 N번의 launch입니다. PR은 이를 step 축을 grid로 펼친 단일 Triton kernel로 바꿉니다.

@triton.jit
def _draft_step_seqlens_kernel(
    seq_lens_ptr, cache_seqlens_ptr, cu_seqlens_k_ptr,
    bs, first_offset, cache_seqlens_stride, cu_seqlens_k_stride,
    BS_BLOCK: tl.constexpr,
):
    # One program per draft step s
    s = tl.program_id(axis=0)
    offs = tl.arange(0, BS_BLOCK)
    mask = offs < bs
    seqlens = tl.load(seq_lens_ptr + offs, mask=mask, other=0).to(tl.int32)
    seqlens = tl.where(mask, seqlens + first_offset + s, 0)
    tl.store(cache_seqlens_ptr + s * cache_seqlens_stride + offs, seqlens, mask=mask)
    tl.store(
        cu_seqlens_k_ptr + s * cu_seqlens_k_stride + 1 + offs,
        tl.cumsum(seqlens, axis=0),
        mask=mask,
    )

def update_draft_step_seqlens(*, seq_lens, cache_seqlens, cu_seqlens_k, bs, steps, first_offset):
    if bs == 0 or steps == 0:
        return
    _draft_step_seqlens_kernel[(steps,)](...,
        BS_BLOCK=triton.next_power_of_2(bs))

포인트는 세 가지입니다.

• grid = (steps,): 각 program이 한 step의 행을 담당하므로 step 간 의존성이 없고 완전 병렬입니다.
• tl.cumsum으로 prefix sum 융합: cache_seqlens를 쓰는 동시에 cu_seqlens_k(누적합)를 같은 레지스터 값에서 계산해 메모리 왕복을 없앴습니다. cu_seqlens_k[s, 0]은 미리 0으로 설정된 값을 그대로 두고 +1 오프셋부터 씁니다.
• tl.where(mask, ..., 0): padding lane을 0으로 만들어 cumsum이 오염되지 않게 합니다.

또한 __init__.py의 kernel 등록 목록에 ("trtllm_mha_graph_metadata", "update_draft_step_seqlens")를 추가해 기존 lazy-load 메커니즘에 편입시켰습니다. 이것이 PR 표의 #4 "Share draft attention metadata"의 핵심이며, TEP8 B8–B128에서 static +7.5%~+12.5%로 측정되었습니다.

2. speculative/eagle.py — chain draft의 acceptance를 한 kernel에

가장 큰 변화는 _chain_accept_kernel입니다. 기존 eagle_sample tail은 verify_tree_greedy → (옵션) simulated acceptance → accept lengths 계산 등 여러 단계였습니다. 새 kernel은 request당 하나의 program이 트리 탐색 전체를 레지스터에서 수행합니다.

b = tl.program_id(0).to(tl.int64)
base = b * WIDTH
cols = tl.arange(0, BLOCK)
if VERIFY:
    # VerifyTreeGreedy over the request's tree, in registers
    last = tl.load(retrieve_index + base).to(tl.int64)
    acc = tl.where(cols == 0, last, -1)
    ...
    for _ in tl.static_range(1, WIDTH):
        active = ~done
        cur = tl.load(retrieve_next_token + base + cur, mask=active, other=-1)
        target = tl.load(target_predict + last).to(tl.int64)
        for _ in tl.static_range(WIDTH):
            searching = active & (cur != -1) & ~matched
            draft = tl.load(candidates + base + cur, mask=searching, other=0)
            hit = searching & (draft.to(tl.int64) == target)
            ...
            sibling = tl.load(retrieve_next_sibling + base + cur,
                              mask=searching & ~hit, other=-1)
            cur = tl.where(searching & ~hit, sibling.to(tl.int64), cur)
        done = done | (active & ~matched)

tl.static_range로 depth/sibling 루프를 컴파일 타임에 unroll하고, done/matched 마스크로 분기 대신 predication을 사용합니다. VERIFY, FINISH, SIMULATE, REAL_DRAFT_TOKENS를 tl.constexpr로 받아 조합별로 특수화된 kernel을 생성하므로, 런타임 분기 비용이 없습니다. docstring이 명시하듯 결과는 eagle_sample과 bitwise 동일을 목표로 합니다. 이렇게 launch 수가 고정된 tail은 CUDA graph로 capture하기에도 적합하며, 이것이 #2 "Capture the post-verify tail"의 기반입니다(TP8 B64 +10.0%, TEP8 B128 +14.1%).

3. _verify_commit_outputs_kernel — 작지만 의미 있는 dtype 일반화

Before:

tl.store(bonus + row, token, valid)
...
bonus = torch.empty((bs,), device=predict.device, dtype=torch.int32)

After:

tl.store(bonus + row, token.to(bonus.dtype.element_ty), valid)
...
def prepare_verify_commit_outputs(..., bonus_dtype=torch.int32):
    bonus = torch.empty((bs,), device=predict.device, dtype=bonus_dtype)

bonus 버퍼를 호출자가 원하는 dtype(예: 다음 graph의 static buffer가 int64인 경우)으로 바로 만들 수 있어 후속 .to() 복사를 제거할 수 있습니다. 기본값 int32에서는 cast가 no-op이라 기본 경로의 동작은 바뀌지 않습니다. 변수명 post → post_length 변경은 가독성 개선입니다.

4. 런타임 측 변경 (diff 축약 구간)

PR 설명에 따르면 나머지 커밋은 다음과 같습니다.

• #1 Vocabulary shard verify: greedy verify는 승리 토큰만 필요하므로 rank-local logits에서 compact key를 만든 뒤 축약된 값만 collective로 reduce. full-vocab all-gather와 임시 텐서, 후속 argmax를 제거. 샘플링·penalty·logprob 요청은 원래 경로로 fallback.
• #5/#6 staging KV scatter 지연 및 pinned/fenced 비동기 readback: TEP8 c272에서 +11.9% / +19.2%.
• #7 PD decode admission을 in-flight forward와 overlap.
• #10 chunked prefill에 맞춘 staging 예약: TTFT P90 1.87s → 0.88s(−53.3%).

환경변수 문서에는 SGLANG_OPT_CUTEDSL_EP_DEFERRED_FINALIZE가 추가되었고, unset이면 마스터 스위치를 따르되 명시적 0/1이 우선하는 계층형 스위치 패턴을 일관되게 적용합니다.

왜 이게 좋은가

측정 결과

• static decode: post-verify tail capture(#1+#2)만으로 TP8 B64 +10.0%, TEP8 B128 +14.1%
• 비동기 staging readback: TEP8 PD c272 P90 +19.2%
• 전체 스택(#47201 + 이 PR + plugin) vs main, TEP8 PD c272: TPS/user P50 87.5 → 301.2 (+244%), TTFT P90 8.02s → 1.99s (−75%). 단, 작성자도 이는 PR 단독 기여가 아니라고 명시합니다.

일반적 교훈

1. 필요한 것만 통신하라: greedy verify에 full logits는 필요 없습니다. reduction을 shard 단위로 먼저 수행하면 collective 대역폭이 vocab 크기에서 독립적이 됩니다.
2. 반복 launch는 grid 축으로 접어라: update_draft_step_seqlens처럼 루프 변수를 program_id로 옮기면 N회 launch가 1회가 됩니다.
3. constexpr 특수화 + predication: 분기 많은 제어 로직도 고정 크기 트리라면 레지스터 안에서 unroll해 하나의 kernel로 만들 수 있고, 이는 CUDA graph capture의 전제가 됩니다.
4. host 동기화는 pinned memory + event fence로 비동기화.
5. 측정 기반으로 덜어내기: 이 PR은 noise 이내 효과의 변경 4개(+3,121 production lines)를 제거해 리뷰 부담을 줄였습니다. 함께 끄면 P50 −3.3%라는 데이터까지 공개해 trade-off를 투명하게 남긴 점은 대형 성능 PR의 좋은 모범입니다.

리뷰 측면에서는 메인테이너(YAMY1234)가 rebase 후 51개 NVIDIA job 통과를 확인했고, rebase로 인한 변경은 #47201 인터페이스에 맞춘 return 값 한 줄뿐이라고 정리되었습니다. 기본 스위치 off에서 main과 동일한 실행을 보장하는 설계 덕분에 회귀 위험을 낮게 유지하며 머지 경로를 확보한 셈입니다.

참고 자료

• CUDA Graphs (PyTorch) — post-verify tail capture에 사용된 CUDA graph 개념
• triton.language.cumsum — draft step seqlens kernel에서 cu_seqlens_k 계산에 사용
• triton.language.static_range — chain accept kernel의 루프 unroll
• triton.language.where — mask 기반 predication
• torch.Tensor.pin_memory — 비동기 staging metadata readback

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