[sglang] FlashInfer MoE autotune의 함정: 모든 토큰이 0이면 튜너가 잘못된 tile을 고른다

PR 링크: sgl-project/sglang#49497 상태: Merged | 변경: +32 / -13

들어가며

SGLang은 서버 시작 시 FlashInfer autotune을 한 번 실행합니다. BaseRunner._flashinfer_autotune이 runner의 static decode buffer 위에서 dummy forward(_dummy_run)를 돌리고, FlashInfer의 MoE tuner는 이 forward에서 라우팅된 row들을 이용해 token bucket(1, 2, 4, …, 64)별로 가장 빠른 kernel tactic([tile_n, config])을 고릅니다.

문제는 그 dummy buffer의 input_ids가 전부 0으로 채워져 있다는 점입니다. 모든 row가 같은 token이므로 router는 모든 row를 같은 top-k expert로 보냅니다. FlashInfer 0.7.3의 make_repeating_tensor_initializer는 profile 크기가 원본보다 작으면 caller의 앞쪽 row를 그대로 반환하기 때문에, tuner는 "소수 expert가 모든 row를 떠안는" 상황에서 tactic 순위를 매깁니다. 이 경우 wide-N tile(예: tile 32)이 유리해 보입니다.

하지만 실제 serving에서는 row가 많은 expert에 흩어져 active expert당 1~2 row 수준이고, 여기서는 좁은 tile(8, 16)이 더 빠릅니다. cold-start에서 잘못 고른 tactic이 서버 수명 내내 쓰인다는 것이 핵심 문제입니다.

코드 분석

1. flashinfer_autotune.py — 새 context manager

핵심은 autotune dummy forward 동안에만 input_ids를 pseudo-random 값으로 바꿔치기하는 context manager입니다.

@contextlib.contextmanager
def autotune_dummy_run_ctx(model_runner: ModelRunner, buffers):
    """Run context of an autotune dummy forward: canary manager and varied token ids."""
    mr = model_runner
    canary_run_ctx = (
        c.with_active_single_forward_manager(0)
        if (c := mr.canary_manager) is not None
        else empty_context()
    )
    with canary_run_ctx:
        input_ids = buffers.input_ids
        vocab_size = mr.model_config.vocab_size
        generator = torch.Generator().manual_seed(get_parallel().attn_dp_rank)
        ids = torch.randint(
            vocab_size, input_ids.shape, generator=generator, dtype=input_ids.dtype
        )
        ...
        saved = input_ids.clone()
        input_ids.copy_(ids)
        try:
            yield
        finally:
            input_ids.copy_(saved)

주목할 설계 포인트:

• [0, vocab_size) 범위: 이 범위의 모든 id는 embedding row를 가집니다. dummy에는 multimodal 입력이 없으므로 special/placeholder token도 텍스트처럼 embed되어 안전합니다.
• seed = attention-DP rank: 같은 batch를 공유하는 TP/CP/PP rank는 동일한 row를 넣어 timing reduction이 같은 조건을 비교하게 하고, attention-DP rank끼리는 서로 다른 row를 넣습니다. DP rank들의 token은 MoE 단계에서 gather/dispatch되므로, 동일하면 같은 route가 dp_size번 반복되어 다시 편향이 생깁니다.
• 고정 seed: 재시작해도 dummy row가 같아 재현성이 유지됩니다.
• CPU generator + copy_: in-place copy로 static buffer의 주소를 보존하고(CUDA graph와 공유되는 buffer), finally에서 원래 값을 복원해 이후 동작에 부작용이 없습니다.

2. cache key 변경

    model_key_parts.append("candidate_selection=v2_default")
+   model_key_parts.append("dummy_tokens=varied")

캐시된 shape는 profiling을 건너뛰므로, constant-token으로 튜닝된 기존 캐시가 그대로 재사용되면 이 PR의 효과가 사라집니다. key에 marker를 추가해 업그레이드 후 딱 한 번 retune되도록 했습니다. 기존 candidate_selection=v2_default와 같은 패턴입니다.

3. base_runner.py / EXTEND pass — 중복 제거와 적용

Before:

canary_run_ctx = (
    c.with_active_single_forward_manager(0)
    if (c := mr.canary_manager) is not None
    else empty_context()
)

def forward_fn():
    self._dummy_run(
        batch_size=batch_size,
        buffers=buffers,
        run_ctx=canary_run_ctx,
    )

After:

def forward_fn():
    self._dummy_run(
        batch_size=batch_size,
        buffers=buffers,
        run_ctx=autotune_dummy_run_ctx(mr, buffers),
    )

decode pass와 maybe_flashinfer_autotune_extend에 복붙되어 있던 canary context 생성 코드가 autotune_dummy_run_ctx 하나로 통합됐습니다. run_ctx 인자로 주입하므로 PP DeepGEMM warmup 같은 다른 _dummy_run 호출자는 영향을 받지 않습니다.

리뷰 논의: FlashInfer #8451 이후에도 의미가 있는가

리뷰어 leejnau는 FlashInfer #8451이 make_repeating_tensor_initializer의 concentrated-route 문제를 직접 고쳤다고 지적하며, 이 PR의 증분 이득을 0.7.4에서 다시 측정해달라고 요청했습니다. 또 최종 코드에서 kill switch(SGLANG_FLASHINFER_AUTOTUNE_VARIED_TOKENS)와 unit test가 제거됐는데 설명이 갱신되지 않았다는 점도 짚었습니다(env var 이름에 ENABLE_/DISABLE_ 동사가 없다는 스타일 리뷰도 있었으나, 결국 변수 자체가 삭제됨).

작성자의 재측정 결과는 흥미롭습니다. #8451은 4 row 이상 bucket의 constant route를 uniform하게 분산된 synthetic route로 교체합니다. 그런데 Kimi(896 experts, top-16)에서 32 row 기준:

route	도달 expert 수
#8451 uniform	~393
실제 router (varied tokens)	253–272

실제 routing은 active expert당 약 1.5배 더 많은 row를 싣습니다. 그 결과 uniform route로 튜닝하면 bucket 32/64에서 tile 8([8,911])을, 실제 router route로 튜닝하면 tile 16([16,391])을 고르며, b ≥ 32에서는 후자가 빠릅니다. 즉 너무 몰린 분포도, 너무 고른 분포도 실제와 다르다는 교훈입니다.

왜 이게 좋은가

FlashInfer 0.7.3 (Kimi-K3.1-Pro, GB300, TP8·DCP8·EP2)

b	constant	varied (평균)	차이
16	17.93	17.34	-3.3%
24	21.54	19.54	-9.3%
32	24.35	21.55	-11.5%
64	27.89	27.46	-1.5%

constant dummy에서는 bucket 32/64에 tile 32가 선택됐고, varied에서는 tile 8/16이 유지됐습니다. downstream fork에서는 prefill+decode 서빙 처리량이 GPU당 +18.6% 향상됐습니다. GSM8K 정확도는 노이즈 범위 내 동일했습니다(같은 kernel의 tactic만 바뀌므로 당연).

FlashInfer 0.7.4 (#8451 포함)

b ≤ 24에서는 차이가 사라지지만 b = 32–64에서 0.9–1.7% 빠름이 유지되었습니다(4 vs 4 cold start, p ≤ 0.035, 두 node pair 모두 같은 방향). 그래서 결론은 "revert하지 않고 유지"였습니다.

일반적 교훈

1. Autotuning은 입력 분포에 민감하다. shape만 맞춘 dummy 입력은 data-dependent kernel(MoE routing, sparse 연산)에서 엉뚱한 최적점을 만듭니다. 튜닝 workload는 실제 분포를 최대한 닮아야 합니다.
2. cold-start 결정은 오래 산다. 시작 시 한 번 정해진 tactic이 서버 수명 내내 쓰이므로, 시작 단계의 작은 편향이 지속적 손실이 됩니다.
3. 캐시 key에 튜닝 조건을 포함하라. 튜닝 방식이 바뀌면 stale cache가 개선을 무력화하므로 marker로 무효화해야 합니다.
4. in-place swap + finally 복원 패턴은 CUDA graph용 static buffer를 건드릴 때 안전한 방식입니다.
5. 분산 환경에서는 어떤 rank끼리 동일한 입력이어야 하고 어떤 rank는 달라야 하는지를 parallelism 구조에 맞춰 seed로 제어해야 합니다.

참고 자료

• torch.randint — dummy token id 생성에 사용된 함수
• torch.Generator — attention-DP rank로 seed를 고정한 CPU generator
• torch.Tensor.copy_ — static buffer를 in-place로 덮어쓰고 복원하는 데 사용
• contextlib.contextmanager — autotune_dummy_run_ctx 구현에 사용된 데코레이터
• FlashInfer #8451 — 리뷰에서 언급된 concentrated-route 문제의 upstream 수정

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