[sglang] Triton JIT 특수화 억제로 디코드 리플레이 첫 패스 컴파일 스톨 85% 줄이기 (SGLang GLM)

PR 링크: sgl-project/sglang#49121 상태: Merged | 변경: +422 / -24

들어가며

LLM 서빙 엔진에서 Triton 커널은 편리하지만 숨은 비용이 있습니다. Triton은 tl.constexpr 인자의 값마다, 그리고 정수 인자가 1이거나 16으로 나누어떨어지는지 같은 성질에 따라 별도 커널을 JIT 컴파일합니다. 배치 크기 bs를 tl.constexpr로 받으면 서버가 처음 보는 batch size가 들어올 때마다 서빙 경로 한가운데서 컴파일이 일어납니다.

이 PR은 GLM-5.3-Flash(4×B200, TP4, C16)에서 나타난 이 문제를 다룹니다. DSA k-pool metadata 커널(decode / verify / draft extend)과 Mamba replay 커널이 리플레이 batch size마다 새 변형을 만들고 있었습니다. 커널 프로브 기준 변형 수는 다음과 같습니다.

• DSA decode / verify / draft extend: [14, 14, 28] → [1, 1, 2]
• Mamba: 49 → 7

그 결과 첫 패스의 일회성 비용은 17.2s → 2.6s(−85%), 첫 패스 streaming throughput은 +17.59% 올랐습니다.

코드 분석

1. scan.py — 하한(floor)이 있는 블록 크기 함수

핵심 아이디어는 공용 헬퍼 하나에 담겨 있습니다.

_MIN_BLOCK_BS = 64
_MIN_BLOCK_EXPANDED = 512

def _metadata_block(n: int, floor: int) -> int:
    """Power-of-two row block covering ``n``, at least ``floor``."""
    block = 1
    while block < n:
        block *= 2
    return max(block, floor)

기존 코드는 triton.next_power_of_2(bs)를 썼습니다. 그래서 BLOCK_BS(constexpr)가 1, 2, 4, …, 64로 바뀔 때마다 변형이 따로 생겼습니다. 하한을 64로 두면 bs ≤ 64 구간 전체가 BLOCK_BS=64 하나로 모입니다. expanded row(bs × next_n)는 더 크므로 하한을 512로 잡았습니다.

2. decode.py — bs를 constexpr에서 런타임 인자로

Before:

@triton.jit(
    do_not_specialize=[
        "page_table_stride_0",
        "real_page_table_stride_0",
        "max_len",
        "num_splits",
    ]
)
def _fused_dsa_decode_metadata_kpool_kernel(
    ...
    bs: tl.constexpr,

After:

@triton.jit(
    do_not_specialize=[
        "page_table_stride_0",
        "real_page_table_stride_0",
        "bs",
        "max_len",
        "num_splits",
    ]
)
def _fused_dsa_decode_metadata_kpool_kernel(
    ...
    bs,

변경은 두 단계입니다.

1. bs: tl.constexpr를 일반 인자 bs로 바꿉니다. 이제 값 자체가 컴파일 키에 들어가지 않습니다.
2. do_not_specialize에 "bs"를 넣습니다. 일반 정수 인자도 Triton은 기본적으로 ==1, %16==0 같은 성질로 특수화하기 때문에, 이것까지 꺼야 bs=1, 16, 32 등이 별도 변형을 만들지 않습니다.

런처 쪽도 함께 바뀝니다.

# Before
block_bs = triton.next_power_of_2(bs)
# After
block_bs = _metadata_block(bs, _MIN_BLOCK_BS)

3. verify.py — constexpr 산술의 제거

verify 커널에는 눈여겨볼 한 줄이 있습니다.

# Before
expanded_size: tl.constexpr = bs * next_n
# After
expanded_size = bs * next_n

bs가 런타임 값이 되면 그 곱도 컴파일 타임 상수가 될 수 없습니다. 따라서 expanded_size도 런타임 스칼라가 되고, 이 값은 마스크 계산에만 쓰입니다. 블록 크기는 아래처럼 하한 함수로 정해집니다.

block_bs = _metadata_block(bs, _MIN_BLOCK_BS)
block_expanded = _metadata_block(expanded_size, _MIN_BLOCK_EXPANDED)

draft_extend.py도 같은 패턴으로 block_expanded = _metadata_block(max_total_len, _MIN_BLOCK_EXPANDED)를 적용합니다. 다만 block_rows = triton.next_power_of_2(max_extend_len)는 그대로 남겨 두었습니다. extend 길이의 범위가 작기 때문입니다.

4. Mamba replay 커널

# Before
@triton.jit
def _fused_replay_state_indices_kernel(...):
...
    BS_UPPER=triton.next_power_of_2(total_bs),

# After
@triton.jit(do_not_specialize=["valid_bs", "total_bs"])
def _fused_replay_state_indices_kernel(...):
...
    # Floored so one compiled variant serves every replay batch size.
    BS_UPPER=max(triton.next_power_of_2(total_bs), 64),

conv-window scatter 커널은 더 흥미롭습니다.

@triton.jit(do_not_specialize=["idx1_ptr", "step1_ptr", "idx2_ptr", "step2_ptr", "n1"])
def _fused_conv_window_scatter_multi_kernel(...)

여기서는 포인터도 특수화 대상에서 뺐습니다. Triton은 포인터 인자의 16바이트 정렬 여부도 컴파일 키에 넣습니다. 그런데 이 커널에는 슬라이스된 텐서의 포인터가 들어오므로 오프셋에 따라 정렬 여부가 달라지고, 그때마다 변형이 갈립니다. 정렬 힌트를 잃는 비용보다 재컴파일을 피하는 이득이 크다고 판단한 것입니다.

리뷰에서의 논의

리뷰어 db-ol의 피드백이 PR의 품질을 크게 끌어올렸습니다.

• 공유 경로 변경 분리: 초기 버전에는 cuda_graph_buffer_registry.py의 small_copy fast path 제거가 포함되어 있었습니다. 리뷰어는 이 경로가 decode/prefill graph runner, eager runner, EAGLE draft runner 전부에 영향을 준다고 지적했습니다. 또 #41166에서 GPU time 11.0→1.4μs와 host dispatch 7.7→57.3μs라는 상충 관계가 측정된 바 있으므로, GLM 한 모델로는 판단할 수 없다고 봤습니다. 결국 이 변경은 PR에서 빠졌습니다.
• 실제 서빙 경로를 테스트하라: draft extend v2 replay는 static_extend_len=True로 실행됩니다. 그런데 STATIC_EXTEND_LEN이 constexpr이므로, 테스트가 dynamic branch만 보면 실제로 컴파일되는 변형을 검증하지 못합니다. 이 지적을 받아 static 케이스가 추가되었습니다.
• 마스크 누락을 잡는 sentinel tail: 하한 64 때문에 bs=1 launch에서는 63개 lane이 마스킹됩니다. 출력 버퍼가 딱 bs 크기라면, 마스크가 빠진 store가 테스트가 읽지 않는 영역에 써도 아무도 알아채지 못합니다. 그래서 출력에 여분 행(TAIL = 3, SENTINEL = -7)을 두고 커널에는 head view만 넘긴 뒤, tail이 그대로인지 확인하도록 바뀌었습니다.
• 마이크로벤치 branch 확인: 처음 측정된 +0.11.1μs 비용은 dynamic branch에서 bs만큼 도는 루프 때문이었습니다. replay가 쓰는 static branch에서는 −3%+5%로 노이즈 수준이었습니다.

왜 이게 좋은가

지표	Base	PR
첫 패스 streaming req/s	2.3272	2.7367 (+17.59%)
일회성 비용(첫 패스 − warm)	17.2s	2.6s
첫 패스 TTFT p95	2.70s	2.07s
warm streaming	2.7576	2.8132
Peak memory	158.90 GiB	158.90 GiB

steady state는 사실상 변화가 없습니다. 이득은 첫 사용 시 컴파일 제거에서만 나옵니다. 그래도 프로덕션에서 이 차이는 큽니다. 새 인스턴스를 띄우거나 autoscaling할 때마다 사용자가 수 초짜리 tail latency를 겪는 일이 사라지기 때문입니다. 출력은 98개 케이스, 451개 버퍼에서 main과 bitwise로 동일했습니다.

일반적 교훈

1. tl.constexpr는 컴파일 키입니다. 값이 다양하게 바뀌는 인자(batch size 등)를 constexpr로 받으면 변형 수가 그 값의 개수만큼 늘어납니다.
2. 일반 정수와 포인터도 기본적으로 특수화됩니다. 정수는 ==1, %16으로, 포인터는 정렬로 구분됩니다. 빈번하게 변하는 인자는 do_not_specialize로 명시하세요.
3. 블록 크기에 하한을 두면 변형이 하나로 모입니다. 작은 bs에서 일부 lane이 낭비되지만, metadata처럼 가벼운 커널에서는 마스킹 비용이 JIT 비용보다 훨씬 작습니다.
4. 마스크에 의존하는 최적화라면 out-of-bounds write 테스트가 필수입니다. sentinel tail 패턴은 저렴하면서도 강력한 방법입니다.
5. 공유 경로 변경은 별도 PR로 분리하고 다른 모델에서도 검증하세요.

참고 자료

• triton.jit — do_not_specialize 등 JIT 특수화 옵션 공식 문서
• triton.language.constexpr — constexpr 인자 공식 문서
• triton.language.arange — 블록 lane 생성에 쓰이는 tl.arange
• CUDA Graphs (PyTorch) — replay 기반 디코드 실행 배경

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