[flashinfer] SM107에서 mm_mxfp8 autotuner가 CUTLASS·CuTe-DSL·cuDNN을 함께 경쟁시키도록 만들기

PR 링크: flashinfer-ai/flashinfer#9020 상태: Merged | 변경: +74 / -21

들어가며

FlashInfer의 mm_mxfp8은 MXFP8(32-element block scale FP8) GEMM을 여러 백엔드(CUTLASS, CuTe-DSL, cuDNN, TRT-LLM, b12x 등)로 실행할 수 있습니다. backend="auto"를 쓰면 heuristic 함수가 후보 목록을 만들고, autotuner(choose_one)가 그중 가장 빠른 runner/tactic을 고릅니다.

문제는 SM107에서 auto가 ["cutlass"] 하나만 반환했다는 점입니다. 후보가 하나뿐이니 autotuner는 CUTLASS tactic끼리만 비교하고, CuTe-DSL이나 cuDNN 커널은 시간 측정조차 하지 않았습니다. CuTe-DSL을 쓰려면 backend="cute-dsl"을 명시해야 했고, 그 경우에도 cuDNN은 측정되지 않았습니다.

그런데 SM107에서는 cuDNN의 커널(2-CTA, B-reuse, 2x2/2x4 cluster)이 일부 대형 prefill shape에서 최고의 CuTe-DSL tactic보다 빠르며, 출력은 bitwise로 동일합니다. 실제로 SGLang은 이 shape들을 cuDNN으로 수동 라우팅(sgl-project/sglang#48971)하고 있었습니다. 이 PR은 그 선택을 autotuner가 스스로 하도록 만듭니다.

코드 분석

1. _heuristic_func_mm_mxfp8: SM107에서만 후보 확장

Before

# don't select trtllm since it requires weight shuffling
if "cutlass" in suitable_backends:
    return ["cutlass"]

After

if "cutlass" in suitable_backends:
    major, minor = get_compute_capability(a.device)
    if (major, minor) != (10, 7):
        return ["cutlass"]
    # On SM107 untuned calls run CUTLASS; the autotuner also ranks CuTe-DSL and cuDNN.
    return [c for c in ("cutlass", "cute-dsl", "cudnn") if c in suitable_backends]

핵심 설계 포인트는 세 가지입니다.

• 순서 보존: CUTLASS가 항상 첫 번째입니다. 튜닝되지 않은 호출이나 cache miss일 때 autotuner는 첫 후보를 fallback으로 쓰므로, 기존 동작(CUTLASS 실행)이 그대로 유지됩니다. 즉 "튜닝할 때만 넓어지는" 확장입니다.
• suitability 필터: 입력 조건상 CuTe-DSL이나 cuDNN이 지원되지 않으면 리스트에서 빠집니다.
• 아키텍처 게이트: SM107 외에는 그대로 ["cutlass"]. 이는 _heuristic_func_mm_fp4가 SM107에서만 cute-dsl을 추가하는 기존 선례를 따른 것입니다.

2. mm_mxfp8: tuning config 선택 조건 완화

Before

tuning_config = (
    _MM_MXFP8_CUTE_DSL_TUNING_CONFIG
    if backends in (["cute-dsl"], ["cutedsl_low_latency"])
    else _MM_MXFP8_TUNING_CONFIG
)

After

tuning_config = (
    _MM_MXFP8_CUTE_DSL_TUNING_CONFIG
    if "cute-dsl" in backends or backends == ["cutedsl_low_latency"]
    else _MM_MXFP8_TUNING_CONFIG
)

이전에는 후보가 정확히 ["cute-dsl"]일 때만 CuTe-DSL용 config(CUDA graph + cold-L2 profiling)를 썼습니다. 이제 후보에 cute-dsl이 섞여 있으면 모든 후보를 같은 방식으로 측정합니다. 서로 다른 측정 방법론(warm L2 vs cold L2, eager vs graph)으로 잰 숫자를 비교하면 랭킹 자체가 왜곡되므로, 공정한 비교를 위한 필수 변경입니다. 게이트 덕분에 실제로 영향을 받는 건 SM107 auto 튜닝뿐입니다.

3. cuDNN runner get_valid_tactics: 실패를 "후보 없음"으로

Before

if self._use_override_shape:
    graph = self._get_override_graph(a, b, out, tactic=0)
else:
    ...
    graph = build_cudnn_gemm_mxfp8_graph(...)
return _cudnn_graph_engine_knob_tactics(graph)

After

try:
    if self._use_override_shape:
        graph = self._get_override_graph(a, b, out, tactic=0)
    else:
        ...
        graph = build_cudnn_gemm_mxfp8_graph(...)
except (cudnn.cudnnGraphNotSupportedError, RuntimeError) as exc:
    # A declined graph leaves the other auto candidates to the tuner.
    logger.debug("cuDNN mm_mxfp8 graph unsupported, no tactics: %s", exc)
    return []
return _cudnn_graph_engine_knob_tactics(graph)

cuDNN이 후보 리스트에 들어오면서, 특정 shape에서 cuDNN이 graph를 거절할 때 예외가 전파되어 auto 튜닝 전체가 실패하는 위험이 생겼습니다. 이제는 빈 tactic 목록을 반환해 다른 후보들의 경쟁을 계속하게 합니다. 이 부분은 리뷰 과정에서 지적되어 추가되었습니다.

4. 테스트

((10, 7), ["cudnn", "cute-dsl", "cutlass"], ["cutlass", "cute-dsl", "cudnn"]),
((10, 0), ["cutlass", "cute-dsl", "cudnn"], ["cutlass"]),
((10, 3), ["cutlass", "cute-dsl", "cudnn"], ["cutlass"]),

get_compute_capability를 monkeypatch해 GPU 없이도 후보 순서와 게이트를 검증하고, declined graph 경로는 override-shape on/off 양쪽에서 []를 확인합니다. 리뷰에서 _use_override_shape를 직접 대입하면 캐시된 runner 상태가 오염된다는 지적이 있어 monkeypatch.setattr로 바뀐 점도 좋은 디테일입니다.

리뷰에서의 핵심 논의

중간 커밋에서는 게이트를 제거해 모든 아키텍처에 후보 확장을 적용했습니다. 리뷰어(leejnau)는 SM100/103에서도 CUTLASS tactic이 다른 tuning config로 측정되어 선택이 바뀔 수 있는데 측정 근거가 없다고 지적했고, cute-dsl JIT 컴파일이 메모리에만 캐시되어 매 프로세스마다 비용을 낸다는 점도 짚었습니다.

작성자는 B300에서 측정한 뒤 게이트를 복원했습니다.

• Autotune 시간: 25 s → 193 s
• wq_b(16384x32768x1280) 1.03–1.04x, wo_b(16384x5120x8192) 1.04–1.08x 느려짐
• 원인은 후보 추가가 아니라 tuning config 전환: CUTLASS만 CuTe-DSL config로 튜닝해도 wq_b가 1.02–1.03x 느려졌음

즉 cold-L2/CUDA-graph 측정으로 고른 tactic이 warm 반복 호출 환경에선 최적이 아닐 수 있다는 흥미로운 결과입니다. 디스크 캐시 지원은 후속 이슈(#9452)로 분리되었습니다.

왜 이게 좋은가

SM107, DeepSeek-V4.1 dense MXFP8 shape 기준(auto / cute-dsl 튜닝 후 시간 비):

M x N x K	auto 선택	ratio
16384 x 25600 x 6144	cuDNN 2x4f_2cta_Breuse	0.76
16384 x 5120 x 8192	cuDNN 2x2f_2cta_Breuse	0.85
16384 x 32768 x 1280	cuDNN 2x4f_2cta_Breuse	0.94
16384 x 4096 x 1280	CuTe-DSL	0.99

SGLang 엔드투엔드 prefill(GPU당 16K 토큰)에서는 앞의 두 shape만 수동 라우팅해도 1.1% 빨라졌는데, 이제 tuner가 이를 자동으로 선택하고 세 번째 shape까지 잡아냅니다.

대가는 튜닝 시간입니다. SM107에서 64 s → 251 s(cute-dsl 명시 시 183 s). 이를 PR 설명에 투명하게 명시한 점, cache shadowing(먼저 cute-dsl로 튜닝된 winner가 auto lookup을 가로채는 기존 동작)을 문서화한 점도 좋습니다.

일반적 교훈

1. 후보 순서로 기본 동작을 보존하라: 확장된 후보 리스트의 첫 원소를 기존 기본값으로 두면, 튜닝하지 않는 사용자에게는 완전히 하위 호환입니다.
2. 같은 잣대로 측정하라: 서로 다른 profiling 방법론의 수치는 비교 불가능합니다. 다만 그 잣대가 실제 워크로드와 맞는지도 검증해야 합니다(B300 회귀 사례).
3. 후보의 실패는 전체 실패가 아니어야 한다: 멀티 백엔드 autotuner에서 개별 runner는 "지원 안 함"을 예외가 아닌 빈 결과로 표현해야 합니다.
4. 측정된 곳에만 적용하라: 이득이 확인된 아키텍처로 범위를 좁히는 게이트는 보수적이지만 안전한 선택입니다.

참고 자료

• cuDNN Frontend (GitHub) — cuDNN MXFP8 graph 구성과 cudnnGraphNotSupportedError를 제공하는 Python frontend
• CUTLASS (GitHub) — CUTLASS 및 CuTe-DSL 백엔드 원본 레포
• pytest monkeypatch — 테스트에서 get_compute_capability 및 runner 속성 패치에 사용

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