[sglang] Blackwell GPU에서 TRT-LLM 커널을 DSA 기본값으로 설정

PR 링크: sgl-project/sglang#21914 상태: Merged | 변경: +2 / -7

들어가며

SGLang의 DeepSeek Sparse Attention(DSA) 기능에서 Blackwell GPU(SM >= 10)의 KV cache dtype과 NSA backend 선택 로직에 dp_size 조건이 걸려 있었다. 이전에는 B200 성능 회귀 이슈(#21291) 때문에 dp_size > 1인 경우에만 fp8_e4m3 KV cache를 사용하고, dp_size == 1인 경우에만 TRT-LLM backend를 선택했다. 이 PR은 해당 조건을 제거하여 Blackwell에서 항상 fp8_e4m3와 TRT-LLM을 기본값으로 사용한다.

핵심 코드 분석

KV Cache dtype 선택 단순화

Before:

if self.kv_cache_dtype == "auto":
    # TODO: Temporarily set default dtype on B200 as bfloat16
    if quantization == "modelopt_fp4" and major >= 10 and self.dp_size > 1:
        self.kv_cache_dtype = "fp8_e4m3"
    else:
        self.kv_cache_dtype = "bfloat16"

After:

if self.kv_cache_dtype == "auto":
    if major >= 10:
        self.kv_cache_dtype = "fp8_e4m3"
    else:
        self.kv_cache_dtype = "bfloat16"

quantization == "modelopt_fp4" 조건과 self.dp_size > 1 조건이 모두 제거되었다. Blackwell이면 무조건 fp8_e4m3을 사용한다.

NSA Backend 선택 단순화

Before:

elif kv_cache_dtype == "fp8_e4m3":
    if self.dp_size == 1 and major >= 10:
        self.nsa_prefill_backend = "trtllm"
        self.nsa_decode_backend = "trtllm"

After:

elif kv_cache_dtype == "fp8_e4m3":
    if major >= 10:
        self.nsa_prefill_backend = "trtllm"
        self.nsa_decode_backend = "trtllm"

dp_size == 1 제약이 제거되어 DP attention 사용 시에도 TRT-LLM backend를 기본 선택한다.

왜 이게 좋은가

• 성능 회귀 해결 완료: B200에서의 성능 이슈가 해결되어 임시 조건문을 제거할 수 있게 됨
• 설정 단순화: 사용자가 dp_size에 따라 다른 동작을 예상하지 않아도 됨
• Blackwell 최적화: fp8_e4m3 + TRT-LLM 조합이 Blackwell에서 최적 성능을 제공

정리

작은 변경이지만, GPU 아키텍처별 기본 동작을 결정하는 핵심 로직이다. Blackwell GPU에서 DSA 관련 성능 이슈가 해결됨에 따라 임시 workaround 조건들을 제거하고 일관된 기본값을 적용한 정리 PR이다.

참고 자료

• sgl-project/sglang#21291 — B200 성능 회귀 이슈
• sgl-project/sglang#21914 — 원본 PR

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