[sglang] SGLang ReplaySSM: GDN 추론 최적화 및 메모리 효율 개선

PR 링크: sgl-project/sglang#28695 상태: Merged | 변경: +1577 / -48

들어가며

최신 LLM 추론 엔진에서 GDN(Gated DeltaNet)과 같은 SSM(State Space Model) 기반 모델은 높은 성능을 보여주지만, Speculative Decoding 과정에서 매 draft token마다 전체 [V, K] 상태를 intermediate_ssm에 기록해야 하는 메모리 오버헤드가 큽니다. 특히 긴 문맥 생성 시 상태 저장 비용은 GPU HBM 대역폭을 크게 점유합니다. 본 PR은 ReplaySSM 기법을 도입하여, 매 단계 상태를 저장하는 대신 per-slot circular ring과 frozen checkpoint를 활용해 메모리 사용량을 획기적으로 줄이고, Closed-loop Exact Fold를 통해 장문 생성 시 발생하는 정확도 저하 문제를 해결합니다.

코드 분석

1. gdn_replayssm_spec_circular_kernel (신규 커널)

이 커널은 기존의 recurrent verify 경로를 대체합니다. 핵심은 상태를 매번 저장하지 않고, 필요한 시점에만 (I + A)^{-1} 변환을 통해 출력을 재구성하는 것입니다.

Before (기존 방식):

# 매 draft token마다 전체 상태를 기록
write_to_intermediate_ssm(full_v_k_state)

After (개선 방식):

# circular ring에 d, k, g만 저장하여 메모리 절감
# 출력은 chunked delta-rule (I + A)^-1로 재구성
phys_c = (b_cache_base + o_c) & (MAX_CACHE_LEN - 1)
# ... (중략) ...
b_d_scaled = (b_d_all * b_replay_decay[None, :]).to(h0.dtype.element_ty)

2. Closed-loop Exact Fold

이전 구현에서는 open-loop fold 방식의 누적 오차로 인해 긴 문맥에서 반복 루프(repetition loop)가 발생하는 문제가 있었습니다. 이를 해결하기 위해 gdn_replayssm_exact_fold_kernel을 통해 원본 입력으로부터 상태를 재연산(replay)하여 정확도를 보장합니다.

# 정확도 유지를 위해 원본 입력으로부터 재연산 수행
# fused_sigmoid_gating_delta_rule_update_kernel과 bit-identical한 연산 순서 유지
beta_s = tl.where(mask_s, 1.0 / (1.0 + tl.exp(-b_s)), 0.0)

왜 이게 좋은가

1. 메모리 효율성: intermediate_ssm 버퍼를 11.48GB에서 0GB로 제거하고, 6개의 링 버퍼(1.81GB)만 사용하여 speculative scratch 메모리를 6.4배 절감했습니다.
2. 정확도 유지: 기존의 open-loop 방식은 생성 길이가 길어질수록 오차가 누적되었으나, Closed-loop Exact Fold를 통해 recurrent baseline과 bit-identical한 결과를 얻어 AIME 벤치마크에서 정확도 패리티를 달성했습니다.
3. 처리량 향상: H20-3e GPU 환경에서 Qwen3.5-122B 모델 기준, 동시성 8에서 출력 처리량(Throughput)이 +30.3% 향상되었으며, TPOT(Time Per Output Token)은 -45.9% 개선되었습니다.

교훈

• Error Path 분리: 출력(Output)은 one-shot으로 계산하여 오차 누적을 방지하고, 상태(State)는 정확한 재연산(Exact Fold)을 통해 누적 오차를 제거하는 설계가 중요합니다.
• 메모리 대역폭 최적화: speculative decoding에서 상태 저장 비용이 병목일 경우, circular buffer를 활용한 재연산 전략이 HBM 대역폭을 확보하는 데 매우 효과적입니다.

리뷰어 피드백 반영

• 정확도 이슈: 초기 구현의 오차 누적 문제를 Closed-loop 방식으로 해결하여 AIME 벤치마크에서 recurrent baseline 대비 동등 이상의 성능을 확인했습니다.
• 메모리 할당: CI에서 발생한 OOM 이슈는 모델 아키텍처(Mamba2 등)와 관련된 것으로 확인되었으며, 본 PR의 ReplaySSM은 선택적(opt-in)으로 활성화하여 안정성을 확보했습니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html
• https://dao-lab.ai/blog/2026/replayssm/

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