본문으로 건너뛰기

[flashinfer] H100에서 D256 Paged Attention을 cuDNN으로 자동 디스패치하기: 실제 causal 작업량 기반 Admission 설계

PR 링크: flashinfer-ai/flashinfer#9451 상태: Merged | 변경: +725 / -12

들어가며

FlashInfer의 BatchPrefillWithPagedKVCacheWrapper는 backend="auto"일 때 native planner가 FA2/FA3 중 하나를 고릅니다. 기존에도 Blackwell(SM100/103)에서는 특정 조건에서 cuDNN으로 넘기는 admission 함수(_blackwell_paged_auto_cudnn)가 있었지만, Hopper에서는 auto가 항상 FA2/FA3에 머물렀습니다.

그런데 H100 80GB HBM3에서 BF16, head_dim 256, HND/page16 causal prefill을 측정해 보면, 작업량이 충분히 큰 경우 cuDNN 9.24 엔진이 FA3보다 확연히 빠릅니다. 예를 들어 B=1, Hq=48, Q=KV=1023에서 graph replay가 100.99µs → 83.67µs로 줄어듭니다. 반대로 작은 workload에서는 cuDNN의 host submission 비용(리뷰어 측정 기준 FA3 32–35µs vs cuDNN 70–74µs)이 이득을 상쇄합니다.

이 PR은 이 "언제 cuDNN이 이기는가"를 선언된 최대 길이가 아닌 실제 causal query-key pair 수로 판단하는 admission 함수를 추가하고, CUDA Graph 수명주기에서 발생하는 table capacity 버그까지 함께 고칩니다.

코드 분석

1. flashinfer/prefill.py — _hopper_paged_auto_cudnn 추가

새 함수는 세 단계로 판단합니다. 먼저 정적 도메인 검사입니다.

if not (
    layout == "HND"
    and q_dtype == kv_dtype == o_dtype == torch.bfloat16
    and head_dim_qk == head_dim_vo == 256
    and 1 <= batch <= 32
    and 16 <= hq <= 48
    and hk > 0
    and hq % hk == 0
    and hq // hk in (1, 2, 4, 8, 16)
    and page_size == 16
    and 65 <= max_q <= 16384
    and max_q <= max_kv <= 32768
    and qo_indptr_host.device.type == "cpu"
    and kv_lens is not None
    and kv_lens.device.type == "cpu"
    and (seq_lens_q is None or seq_lens_q.device.type == "cpu")
):
    return False

주목할 점은 길이 정보가 이미 CPU에 있을 때만 진입한다는 것입니다. device 텐서라면 .tolist()가 동기화(readback)를 유발하므로, 이를 원천 차단하는 fail-closed 설계입니다. max_q >= 65는 기존 short-query FA2 도메인(Q≤64)을 건드리지 않기 위함입니다.

다음은 선언값과 실제값의 일관성 검사입니다.

bounds = qo_indptr_host.tolist()
queries = [end - start for start, end in zip(bounds, bounds[1:], strict=False)]
lengths = kv_lens.flatten().tolist()
if not (
    len(queries) == len(lengths) == batch
    and bounds[0] == 0
    and bounds[-1] == total_q
    and max(queries) == max_q
    and max(lengths) == max_kv
    and all(0 < q <= kv for q, kv in zip(queries, lengths, strict=True))
    and (seq_lens_q is None or seq_lens_q.flatten().tolist() == queries)
):
    return False

메타데이터가 조금이라도 어긋나면 cuDNN을 선택하지 않습니다. 잘못된 선언으로 cuDNN graph를 만들면 정확성 문제가 될 수 있으므로, 의심스러울 땐 검증된 native 경로로 남는 보수적 접근입니다.

핵심은 작업량 하한입니다.

# A product of declared maxima would admit nearly empty ragged batches.
paired_work = hq * sum(
    q * (2 * kv - q + 1) // 2 for q, kv in zip(queries, lengths, strict=True)
)
if paired_work < 1 << 24:
    return False

Causal mask에서 query 길이 q, KV 길이 kv(q≤kv, 끝 정렬)일 때 i번째 query는 kv - q + i + 1개의 key를 봅니다. 이를 합하면 등차수열 합 q * (2*kv - q + 1) / 2가 됩니다. 이 값을 요청별로 계산하고 head 수를 곱해 실제 attention 연산량에 비례하는 지표를 얻습니다. 만약 batch * max_q * max_kv 같은 선언값 곱을 썼다면, 한 요청만 길고 나머지는 거의 비어 있는 ragged batch가 부당하게 admit되었을 것입니다. 2**24는 PR 설명대로 "아키텍처 정리"가 아닌 독립 입력으로 검증된 경험적 임계값입니다.

마지막으로 하드웨어 식별입니다.

return (
    get_compute_capability(device) == (9, 0)
    and get_device_sm_count(device) == 132
    and get_device_name(device) == "NVIDIA H100 80GB HBM3"
    and _cudnn_supports_paged_auto()
)

SM 9.0이라도 H20, PCIe H100, GH200은 제외됩니다. 비용이 큰 device 쿼리를 가장 마지막에 두어 단락 평가로 대부분의 호출에서 건너뜁니다.

2. plan() — provider별 admission 분기

# Before
if (
    planned_backend == "fa2"
    and (not self.is_cuda_graph_enabled or self._plan_info is None)
    ...
    and _blackwell_paged_auto_cudnn(

# After
cudnn_admission = (
    _blackwell_paged_auto_cudnn
    if planned_backend == "fa2"
    else _hopper_paged_auto_cudnn
)
if (
    planned_backend in ("fa2", "fa3")
    and (not self.is_cuda_graph_enabled or self._plan_info is None)
    ...
    and cudnn_admission(

native planner가 FA2를 고르면 기존 Blackwell 함수, FA3를 고르면 새 Hopper 함수로 판단합니다. 즉 "native selector가 먼저 FA3를 고른다"는 전제를 코드 구조로 강제합니다. self._plan_info is None 조건 덕분에, 이미 native schedule을 소유한 graph-mode wrapper는 provider를 바꾸지 않아 retained capture가 유효하게 유지됩니다.

또한 Blackwell용 bucket 선언은 이제 SM100/103에서만 적용됩니다.

if (
    self._requested_backend == "auto"
    and self._backend == "cudnn"
    and get_compute_capability(self.device) in ((10, 0), (10, 3))
):

Hopper는 explicit cuDNN과 동일한 exact declaration을 쓰므로, 검증되지 않은 bucket 전략이 Hopper로 새지 않습니다.

3. Owned table capacity 수정 (최종 fix)

# Before
if self._cudnn_paged_bounds is not None:
    max_pages = max(max_pages, self._cudnn_paged_bounds[1] // page_size)
    if self.is_cuda_graph_enabled:
        index_capacity = self._paged_kv_indices_buf.numel()
        admitted_pages = min(index_capacity, 32768 // page_size)
        max_pages = max(max_pages, index_capacity, 1 << (admitted_pages - 1).bit_length())

# After
if self._cudnn_paged_bounds is not None:
    max_pages = max(max_pages, self._cudnn_paged_bounds[1] // page_size)
if self.is_cuda_graph_enabled:
    index_capacity = self._paged_kv_indices_buf.numel()
    max_pages = max(max_pages, index_capacity)
    if self._cudnn_paged_bounds is not None:
        admitted_pages = min(index_capacity, 32768 // page_size)
        max_pages = max(max_pages, 1 << (admitted_pages - 1).bit_length())

기존엔 index capacity 예약이 Blackwell bucket 분기 안에만 있었습니다. Hopper는 exact declaration이라 _cudnn_paged_bounds가 None이므로, Q65/KV8112로 첫 plan 시 507 페이지만 할당한 뒤 KV8192로 재plan하면 captured cuDNN block table cannot grow 에러가 났습니다. 이제 graph mode라면 선언 방식과 무관하게 전체 index 버퍼 용량을 예약합니다. 이 버그는 red/green GPU 테스트로 재현·검증되었습니다.

리뷰에서 나온 논의

YangXu1990uiuc는 이 capacity 변경의 부작용을 정량적으로 지적했습니다. 각 row가 전체 batch의 index capacity를 예약하므로 uniform 요청에서 저장량이 B * pages에서 B² * pages로 늘어납니다. explicit backend="cudnn" 사용자에게도 영향이 갑니다.

Batch Owned table Host plan()
16 32 KiB → 512 KiB 127 → 214 µs
256 512 KiB → 128 MiB 170 → 1210 µs

또한 rectangular-copy fast path도 놓치게 됩니다. 리뷰어는 이를 correctness가 아닌 non-blocking 자원 이슈로 분류하고, per-request capacity 상한이나 growth opt-in 같은 후속 개선을 제안했습니다. caller-owned block_tables를 넘기면 회피 가능합니다.

같은 리뷰어의 측정에서 (4, 16/4, 65, 4096)은 replay는 144→114µs로 빨라졌지만 completed eager는 177→180µs로 오히려 비슷하거나 느렸습니다. cuDNN의 host submission 비용 때문이며, 작업량 하한이 왜 필요한지 보여줍니다.

0z5a는 GH200에서도 cuDNN/FA3 비율이 H100과 거의 같음(0.78~0.83)을 보였지만, 이 PR은 의도적으로 이름 기반 제외를 유지해 별도 qualification을 요구합니다.

왜 이게 좋은가

  • 성능: 변경된 173개 route에서 warm graph replay 1.260x, completed eager 1.266x 개선. 변경되지 않은 147개 FA2/FA3 대조군은 [0.998, 1.002] 구간으로 회귀 없음.
  • 비용 모델의 정확성: 선언값이 아닌 실제 causal pair 수로 판단해 ragged batch 오판을 막음.
  • Zero readback: CPU에 있는 길이만 사용하고, 없으면 fail-closed.
  • Graph 수명주기 안전성: 기존 capture를 가진 wrapper는 provider를 바꾸지 않음.

일반적 교훈은 세 가지입니다. 첫째, backend dispatch는 "어느 커널이 빠른가"가 아니라 "이 workload에서 고정 오버헤드를 상각할 만큼 일이 큰가"의 문제입니다. 둘째, 측정된 하드웨어 밖으로 결론을 일반화하지 않는 좁은 admission이 안전합니다. 셋째, 버그 수정(capacity)이 다른 사용자 경로에 비용을 전가할 수 있으므로, 리뷰어처럼 메모리·plan 시간을 실측하는 습관이 중요합니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글