본문으로 건너뛰기

[SGLang] EAGLE v2: 개선된 드래프트 알고리즘

들어가며

EAGLE v2는 SGLang에서 "overlap schedule"이 활성화될 때 사용되는 개선된 speculative decoding 워커다. V1(EAGLEWorker)이 드래프트-검증을 순차 실행하는 반면, V2(EAGLEWorkerV2)는 드래프트와 타겟 검증을 비동기로 오버랩하여 파이프라인 효율을 높인다. 핵심 차이는 EagleDraftInputV2Mixin과 EagleVerifyInputV2Mixin에 있으며, KV 캐시 사전 할당, 비동기 verify, extend-to-fill 전략 등이 추가되었다.

구조도

EAGLE V1 (동기):
  Draft ──▶ Verify ──▶ Draft ──▶ Verify ──▶ ...
  ████████  ████████  ████████  ████████

EAGLE V2 (오버랩):
  Draft₁ ──▶ Draft₂ ──▶ Draft₃ ──▶ ...
  ████████   ████████   ████████
       Verify₁ ──▶ Verify₂ ──▶ ...
       ████████   ████████
  └─ 드래프트와 검증이 파이프라인으로 겹침 ─┘

핵심 코드 분석

1. EagleDraftWorker: V2 전용 워커

python/sglang/srt/speculative/eagle_worker_v2.py에서 EagleDraftWorker는 BaseDraftWorker를 구현한다.

class EagleDraftWorker(BaseDraftWorker):
    def __init__(self, server_args, gpu_id, tp_rank, dp_rank,
                 moe_ep_rank, attn_cp_rank, moe_dp_rank,
                 nccl_port, target_worker):
        self.draft_worker = TpModelWorker(
            server_args=server_args, gpu_id=gpu_id,
            is_draft_worker=True,
            req_to_token_pool=self.req_to_token_pool,
            token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
            memory_pool_config=target_worker.model_runner.memory_pool_config,
        )
        self.draft_runner = self.draft_worker.model_runner

V2에서는 드래프트 워커가 TpModelWorker를 내부에 별도로 생성하여, 타겟 워커와 독립적으로 forward를 실행할 수 있다.

2. KV 캐시 사전 할당: prepare_for_decode

V2의 EagleDraftInputV2Mixin은 디코드 시작 전에 KV 캐시를 미리 할당한다.

def prepare_for_decode(self: EagleDraftInput, batch: ScheduleBatch):
    alloc_len_per_decode = get_alloc_len_per_decode()
    for r in batch.reqs:
        x = r.kv_committed_len + 2 * alloc_len_per_decode - r.kv_allocated_len
        cur_kv_lens_cpu.append(r.kv_allocated_len)
        nxt_kv_lens_cpu.append(r.kv_allocated_len + x)
        r.kv_allocated_len += x

    if page_size == 1:
        out_cache_loc = alloc_token_slots(batch.tree_cache, num_needed_tokens)
    else:
        out_cache_loc = alloc_paged_token_slots_extend(...)

2 * alloc_len_per_decode만큼 초과 할당하여, 드래프트와 검증이 동시에 진행될 때 캐시 부족을 방지한다. V1에서는 매 step마다 할당/해제를 반복하지만, V2는 사전 할당으로 동기화 오버헤드를 줄인다.

3. 드래프트 캐시 위치 할당: Triton 커널

V2 전용 Triton 커널이 드래프트 캐시 위치를 계산한다.

@triton.jit
def assign_draft_cache_locs_page_size_1(
    req_pool_indices, req_to_token, seq_lens, out_cache_loc,
    pool_len: tl.constexpr, topk: tl.constexpr,
    speculative_num_steps: tl.constexpr,
):
    pid = tl.program_id(axis=0)
    copy_len = topk * speculative_num_steps
    out_cache_ptr = out_cache_loc + pid * topk * speculative_num_steps
    kv_start = tl.load(seq_lens + pid)
    token_pool = req_to_token + tl.load(req_pool_indices + pid) * pool_len
    for i in range(num_loop):
        copy_offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE
        data = tl.load(token_pool + kv_start + copy_offset, mask=...)
        tl.store(out_cache_ptr + copy_offset, data, mask=...)

이 커널은 page_size == 1일 때 각 요청의 KV 캐시 위치를 req_to_token 풀에서 out_cache_loc로 복사한다.

4. 비동기 검증 준비: prepare_for_v2_verify

def prepare_for_v2_verify(self: EagleVerifyInput, req_to_token_pool, batch, target_worker):
    bs = len(batch.req_pool_indices)
    batch.input_ids = self.draft_token
    batch.out_cache_loc = assign_extend_cache_locs_func(
        req_pool_indices=batch.req_pool_indices,
        req_to_token=req_to_token_pool.req_to_token,
        start_offset=batch.seq_lens,
        end_offset=batch.seq_lens + self.draft_token_num,
        batch_size=bs, draft_token_num=self.draft_token_num,
        device=device,
    )
    batch.forward_mode = ForwardMode.TARGET_VERIFY
    verify_forward_batch = ForwardBatch.init_new(batch, target_worker.model_runner)

V2에서는 검증 forward batch를 별도로 구성하여 타겟 워커에 전달한다. CUDA graph 사용 가능 여부도 여기서 판단한다.

5. 검증 결과 처리: sample 메서드

V2의 sample 메서드는 V1의 verify와 유사하지만, 결과를 predict, accept_length, accept_index 텐서로 반환하여 오버랩 스케줄러가 후처리할 수 있게 한다.

def sample(self: EagleVerifyInput, batch, logits_output, vocab_mask=None):
    if sampling_info.is_all_greedy or _is_npu or _is_hip:
        predict, accept_index, accept_length = verify_tree_greedy_func(...)
    else:
        tree_speculative_sampling_target_only(...)
    # V2에서는 bonus token을 포함
    accept_length.add_(1)
    return predict, accept_length, accept_index

accept_length.add_(1)로 bonus token(검증 단계에서 추가로 확정되는 토큰)을 포함시킨다.

6. Extend-to-Fill: 드래프트 KV 캐시 채우기

V2는 검증 후 수락된 토큰으로 드래프트 KV 캐시를 갱신하는 extend 단계를 추가한다.

def prepare_for_extend_to_fill_draft_kvcache(self, batch, predict, num_draft_tokens, ...):
    batch.input_ids = predict
    batch.seq_lens = batch.seq_lens + num_draft_tokens
    batch.forward_mode = ForwardMode.DRAFT_EXTEND_V2
    forward_batch = ForwardBatch.init_new(batch, draft_model_runner)

ForwardMode.DRAFT_EXTEND_V2는 V2 전용 forward mode로, 검증된 토큰을 드래프트 모델의 KV 캐시에 반영한다.

V1 vs V2 비교

항목 EAGLE V1 EAGLE V2
실행 방식 동기 (Draft → Verify → Draft) 비동기 오버랩
KV 캐시 할당 매 step 동적 할당 사전 할당 (2x buffer)
Forward Mode DECODE, TARGET_VERIFY + DRAFT_EXTEND_V2
워커 구조 EAGLEWorker (TpModelWorker 상속) EagleDraftWorker + EAGLEWorkerV2
검증 결과 EagleVerifyOutput dataclass (predict, accept_length, accept_index) tuple

설계 근거

오버랩 스케줄링: GPU utilization을 높이려면 드래프트와 검증이 서로의 완료를 기다리지 않아야 한다. V2는 드래프트 forward와 타겟 verify를 CUDA stream 수준에서 겹쳐 실행한다.

사전 할당 전략: 동적 할당은 cudaMalloc 동기화를 유발할 수 있다. 2 * alloc_len_per_decode 초과 할당으로 대부분의 경우 재할당 없이 진행할 수 있다.

Plan Stream 분리: _get_plan_stream으로 attention metadata 계산을 별도 CUDA stream에서 실행하여 메인 compute stream과 오버랩한다.

관련 포스트

참고

댓글

관련 포스트

SGLang 의 다른글