[sglang] sgl-kernel 레거시 경로 대규모 제거: AOT에서 JIT로의 전환 완성

PR 링크: sgl-project/sglang#21528 상태: Merged | 변경: +14 / -1708

들어가며

SGLang은 GPU 커널을 두 가지 방식으로 제공해왔습니다: AOT(Ahead-of-Time) 빌드인 sgl-kernel 패키지와, 실행 시점에 컴파일하는 JIT(Just-in-Time) 커널입니다. JIT 커널이 성숙해지면서 AOT 빌드의 일부 커널이 불필요해졌고, 이번 PR은 그 레거시 코드를 대량 제거합니다.

핵심 코드 분석

1. CUDA 소스 파일 삭제

CMakeLists.txt에서 더 이상 빌드하지 않는 소스 파일을 제거:

Before:

set(SOURCES
    "csrc/attention/cascade.cu"
    "csrc/attention/cutlass_mla_kernel.cu"
    "csrc/elementwise/cast.cu"
    "csrc/elementwise/concat_mla.cu"
    # ...

After:

set(SOURCES
    "csrc/attention/cutlass_mla_kernel.cu"
    "csrc/elementwise/concat_mla.cu"
    # ...

cascade.cu(55줄)와 cast.cu가 삭제되었습니다. 이들은 FlashInfer와 JIT 커널로 각각 대체되었습니다.

2. 벤치마크에서 AOT 참조 제거

Before (bench_cast.py):

from sgl_kernel import downcast_fp8 as downcast_fp8_aot

LINE_VALS = ["aot", "jit"]
LINE_NAMES = ["AOT (sgl-kernel)", "JIT (cast.cuh, 256 threads, 2D grid)"]

if provider == "aot":
    fn = lambda: downcast_fp8_aot(k, v, k_out, v_out, k_scale, v_scale, loc)
else:
    fn = lambda: downcast_fp8_jit(k, v, k_out, v_out, k_scale, v_scale, loc)

After:

LINE_VALS = ["jit"]
LINE_NAMES = ["JIT (cast.cuh, 256 threads, 2D grid)"]

fn = lambda: downcast_fp8_jit(k, v, k_out, v_out, k_scale, v_scale, loc)

3. 테스트에서 레거시 기능 제거

test_renorm.py에서 test_top_k_mask_logits 테스트가 삭제되고, bench_renorm.py에서 관련 벤치마크도 제거되었습니다.

왜 이게 좋은가

1. 빌드 단순화: AOT 빌드에서 불필요한 CUDA 파일을 제거하여 빌드 시간이 단축됩니다.
2. 유지보수 부담 감소: 1,708줄의 중복 코드 제거로, 두 가지 구현을 동시에 관리할 필요가 없어집니다.
3. JIT 전환 완성: AOT와 JIT 중 하나만 유지함으로써, 커널 업데이트 시 한 곳만 수정하면 됩니다.

정리

1,700줄 이상의 대규모 삭제지만, 본질은 AOT에서 JIT로의 전환 완성입니다. 벤치마크와 테스트에서도 AOT 참조를 깔끔하게 제거하여, 더 이상 sgl_kernel import에 의존하지 않습니다.

참고 자료

• sgl-project/sglang#21528
• sgl-kernel/CMakeLists.txt

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