[sglang] [성능 최적화] 불필요한 Tree Mask Fill 제거를 통한 Speculative Decoding 가속화

2026년 7월 30일수정: 2026년 7월 30일

PR 링크: sgl-project/sglang#32886 상태: Merged | 변경: +130 / -10

들어가며

대규모 언어 모델(LLM)의 추론 속도를 높이기 위한 기술 중 하나인 Speculative Decoding은 여러 개의 토큰 후보를 한 번에 생성(Drafting)하고, 이를 타겟 모델이 한 번에 검증(Verification)하는 방식을 취합니다. 이 과정에서 SGLang의 EAGLE과 같은 알고리즘은 'Tree-structured' 검증을 사용하는데, 이때 각 토큰 간의 인과 관계를 정의하기 위해 tree_mask라는 거대한 불리언 행렬을 생성합니다.

기존 구현에서는 매 디코딩 단계마다 이 tree_mask 버퍼 전체를 True로 초기화(memset)하는 과정이 포함되어 있었습니다. 하지만 컨텍스트 길이가 길어질수록 이 버퍼의 크기는 수백 MB에 달하게 되며, 매 단계 이를 초기화하는 것은 상당한 메모리 대역폭 낭비와 지연 시간을 초래합니다.

이번 PR([Perf] Skip the target-verify tree mask fill)은 특정 조건에서 백엔드가 이 마스크를 읽지 않는다는 점에 착안하여, 불필요한 초기화 과정을 생략함으로써 성능을 최적화했습니다. 시니어 엔지니어의 관점에서 이 변경사항이 왜 중요한지 분석해 보겠습니다.

코드 분석: 무엇이 바뀌었는가?

1. Backend Interface 확장 (`base_attn_backend.py` 외)

먼저, 각 어텐션 백엔드가 custom_mask를 실제로 읽는지 여부를 알려줄 수 있는 훅(Hook)을 추가했습니다. 기본값은 하위 호환성을 위해 True로 설정되었습니다.

# python/sglang/srt/layers/attention/base_attn_backend.py

def target_verify_reads_custom_mask(self) -> bool:
    """Whether target-verify attention reads spec_info.custom_mask at all.

    When False, build_tree_kernel_efficient skips the full-buffer prefix
    fill (max_num_tokens x max_context_len bool memset per verify step).
    """
    return True

이후 FlashAttention 백엔드와 DeepSeek V4 백엔드에서 이를 최적화합니다. 특히 FlashAttention의 경우 topk <= 1일 때는 마스크를 참조하지 않는다는 특성을 이용합니다.

# python/sglang/srt/layers/attention/flashattention_backend.py

def target_verify_reads_custom_mask(self) -> bool:
    # topk<=1 verify never extracts from custom_mask
    return self.topk > 1

2. 핵심 최적화 로직 (`eagle_utils.py`)

가장 중요한 변경은 build_tree_kernel_efficient 함수에 있습니다. 기존에는 tree_mask_mode가 FULL_MASK일 경우 무조건 전체를 fill_(True) 했으나, 이제는 fill_prefix_mask 플래그에 따라 이를 제어합니다.

Before:

elif tree_mask_mode == TreeMaskMode.FULL_MASK:
    tree_mask.fill_(True)

After:

elif tree_mask_mode == TreeMaskMode.FULL_MASK:
    # Only the [0, seq_len) prefix columns depend on this fill;
    # the kernel below writes every tree cell itself.
    if fill_prefix_mask:
        tree_mask.fill_(True)

또한, 버퍼를 새로 할당해야 하는 경우에도 torch.full 대신 초기화되지 않은 메모리를 할당하는 torch.empty를 사용하여 성능을 극대화했습니다.

Before:

elif tree_mask_mode == TreeMaskMode.FULL_MASK:
    tree_mask = torch.full(
        (seq_lens_sum * num_verify_tokens + num_verify_tokens * num_verify_tokens * bs,),
        True,
        device=device,
    )

After:

elif tree_mask_mode == TreeMaskMode.FULL_MASK:
    mask_shape = (seq_lens_sum * num_verify_tokens + num_verify_tokens * num_verify_tokens * bs,)
    tree_mask = (
        torch.full(mask_shape, True, dtype=torch.bool, device=device)
        if fill_prefix_mask
        else torch.empty(mask_shape, dtype=torch.bool, device=device)
    )

왜 이게 좋은 최적화인가?

1. 메모리 대역폭(Memory Bandwidth) 절약

현대 GPU 연산에서 성능 병목은 종종 연산 능력(FLOPS)이 아니라 메모리 대역폭에서 발생합니다. max_num_tokens가 512이고 max_context_len이 32k인 경우, bool 타입 마스크는 약 16MB 정도를 차지합니다. 하지만 배치 사이즈가 커지고 여러 레이어를 거치며 매 스텝마다 이를 memset 하는 비용은 무시할 수 없습니다. 특히 Speculative Decoding의 Verification 단계는 매우 빈번하게 일어나므로, 이 작은 절약이 모여 전체 처리량(Throughput) 향상으로 이어집니다.

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

PR Analysis 의 다른글

이전글 [vllm] vLLM DeepSeek-V4 성능 최적화: 불필요한 torch.full 커널 제거로 1.88배 속도 향상
현재글 : [sglang] [성능 최적화] 불필요한 Tree Mask Fill 제거를 통한 Speculative Decoding 가속화
다음글 [vllm] vLLM DeepSeek-V4 성능 최적화: Eager Break 시 Workspace 재사용 개선

[sglang] [성능 최적화] 불필요한 Tree Mask Fill 제거를 통한 Speculative Decoding 가속화

들어가며

코드 분석: 무엇이 바뀌었는가?

1. Backend Interface 확장 (`base_attn_backend.py` 외)

2. 핵심 최적화 로직 (`eagle_utils.py`)

왜 이게 좋은 최적화인가?

1. 메모리 대역폭(Memory Bandwidth) 절약

2. 불필요한 연산의 조건부 제거

3. 테스트를 통한 안정성 검증

리뷰어 피드백 분석

마치며

참고 자료

댓글

관련 포스트

PR Analysis 의 다른글