[flashinfer] FP8 가중치를 커널 안에서 풀어낸다: FlashInfer의 W8A16 GEMM(mm_bf16_fp8) 분석

PR 링크: flashinfer-ai/flashinfer#7495 상태: Merged | 변경: +8610 / -0

들어가며

FP8 e4m3 가중치와 per-tensor FP32 scale 하나로 양자화된 모델(TRT-LLM의 FP8QDQ loader가 만드는 형태)을 BF16/FP16 activation으로 돌리려면, 지금까지는 대개 cast-copy 경로를 썼습니다.

# 기존 방식 (TRT-LLM W8A16FP8LinearMethod 형태)
w16 = weight.to(dtype).mul_(scale.to(dtype))  # 16-bit 가중치를 통째로 materialize
out = F.linear(act, w16)

이 방식에는 두 가지 비용이 있습니다.

1. 메모리 트래픽: N*K 크기의 FP8 가중치를 읽어 2배 크기의 16-bit 사본을 쓰고, GEMM이 그것을 다시 읽습니다.
2. 정밀도: scale 자체가 activation dtype으로 반올림되고, weight * scale 결과도 한 번 더 반올림된 뒤 GEMM에 들어갑니다.

이 PR은 flashinfer.gemm.mm_bf16_fp8 / mm_fp16_fp8을 추가해, FP8 가중치를 커널 내부(레지스터/TMEM)에서 변환하고 FP32 scale은 FP32 accumulator에 적용합니다. 16-bit 가중치 사본은 어디에도 생기지 않습니다.

코드 분석

1. 구현 전략: CUTLASS C++ → CuTe-DSL

초기 리비전은 CUTLASS C++ 포크(include/flashinfer/gemm/w8a16_sm100/kernel.hpp 등)였지만, 최종본은 이를 전부 CuTe-DSL(Python, nvidia-cutlass-dsl)로 재작성했습니다. 계기는 리뷰어 jimmyzho의 지적이었습니다.

CUTLASS 4.8.0으로 빌드하면 두 곳에서 실패합니다. 4.8에서 one-argument TileScheduler::can_implement가 제거됐는데 포크된 커널이 여전히 호출하고, SM100 epilogue builder가 CtaN=160을 static-assert로 막습니다.

CuTe-DSL로 옮기면서 포크 유지 부담이 사라졌고, 256x160 tile은 DSL 커널 자체 epilogue((128, 64) TMA-store subtile)로 처리됩니다. 아키텍처·tile·dtype·bias 조합별로 첫 호출 시 JIT 컴파일되므로 AOT 패키지 크기나 nvcc 빌드에 영향이 없습니다. CI에서 한때 Ninja build failed로 B200/GB200 전체가 깨졌던 것도 이 전환 이후 사라졌습니다.

• Hopper (sm_90a): warp-specialized WGMMA mainloop, fp8 fragment를 레지스터에서 변환. cooperative 128-row / ping-pong 64-row tile, SM을 채우지 못하는 grid용 fused split-K fix-up.
• Blackwell (sm_100a, sm_103a): 2-SM/1-SM tcgen05 tile(256x128, 256x160, 256x192, 128x128), 전용 transform warpgroup이 가중치를 TMEM으로 변환.

2. 공개 API: mm_* 규약에 맞추기

처음 이름은 w8a16_fp8_gemm(..., algo=...)이었습니다. 리뷰어는 형제 API인 mm_bf16_fp4(a, b, b_descale, alpha, ...), mm_fp8(a, b, alpha, out_dtype, out, backend)와 일관성을 요구했고, 최종 형태는 다음과 같습니다.

# After (d958ff315)
mm_bf16_fp8(a, b, alpha, *, bias, out_dtype, out, backend, tactic, workspace)
mm_fp16_fp8(a, b, alpha, *, bias, out_dtype, out, backend, tactic, workspace)
mm_bf16_fp8_workspace_bytes(...)

algo 문자열은 tactic으로 이름만 바뀌었습니다. autotuner가 없으므로 index가 아닌 이름을 유지했다는 설명이 붙어 있습니다. 출시 후 rename은 API break이므로 merge 전에 정리한 것은 좋은 판단입니다.

3. 리뷰로 잡힌 커널 버그들

(a) 정렬 가정 오류 — bias가 없을 때 stand-in 포인터(weight_scale)를 16-byte aligned로 선언해, scale 주소가 % 16 == 4이면 make_ptr가 AssertionError를 냈습니다. 기존 테스트는 reference 호출이 먼저 커널을 컴파일해 둔 탓에 통과했습니다.

# After (47d77b0ac): 세 spec 모두 공유
bias_align = 16 if has_bias else 4

테스트도 _KERNEL_CACHE를 비우고 FLASHINFER_CUTE_DSL_DISABLE_CACHE=1로 강제 재컴파일하도록 고쳤습니다. 캐시가 테스트 실패를 가리는 전형적인 사례입니다.

(b) magic-number 나눗셈 overflow — d == 1이면 magic이 2^40이라 x >= 2^23에서 x * magic이 Int64를 넘어 illegal instruction이 발생했습니다. d == 1 특수처리 대신 x_max > INT64_MAX // magic이면 0(fallback)을 반환하는 일반 조건으로 고쳤습니다(a96bfc02e).

(c) wgmma.fence 누락 — convert가 다음 k-block의 A fragment 레지스터를 쓰고 바로 cute.gemm이 읽는데 fence가 없었습니다(PTX에 16 wgmma에 8 fence). ptxas가 WARPGROUP.ARRIVE를 자동 삽입해 우연히 맞았을 뿐이므로, convert마다 warpgroup.fence()를 넣고 루프 뒤의 무의미한 fence는 제거했습니다(ccf6f1183).

(d) split-K tail의 deadlock 가능성 — tail tile의 모든 slice(최대 16)가 동시에 resident여야 완료되는 구조라, 다른 커널이 SM을 점유하면 9.5초씩 멈췄습니다. 저자는 "마지막 도착 slice가 혼자 reduce" 안이 16-way tail을 17→45 us로 느리게 만든다는 것을 측정한 뒤, 50 us bounded wait + CAS 기반 withdraw를 택했습니다. 시간 초과한 slice가 비트를 세워 빠지고, 남은 slice들이 column을 나눠 맡되 합산은 slice 순서대로 하므로 결과가 bitwise 동일합니다. 재현 시나리오는 12 s → 0.78 ms.

4. Blackwell tile walk order

Nsight Compute 분석에서 persistent scheduler가 activation tile을 빠르게 순회해 더 큰 bf16 activation을 L2에 붙잡아 두어야 하는 상황이 드러났습니다. n < 2 m(fp8 가중치가 더 작은 operand)일 때 가중치 row를 빠르게 순회하도록 바꾼 결과, M=8192, N=10240, K=5120에서 L2 hit rate 59%→84%, DRAM read 1.27 GB→0.32 GB였습니다.

M	N	K	old	new	gain
32768	3200	5120	853 us	683 us	19.9%
16384	12800	5120	1630 us	1400 us	14.1%
8192	10240	5120	578 us	535 us	7.4%

5. 벤치마크: benchmarks/bench_mm_bf16_fp8.py

측정 방법론도 꼼꼼합니다. L2 flush와 측정 대상을 별도 CUDA graph로 캡처합니다.

g_reset = torch.cuda.CUDAGraph()
with torch.cuda.graph(g_reset):
    if args.l2_flush:
        _l2_flush.fill_(0.0)
...
g_run = torch.cuda.CUDAGraph()
with torch.cuda.graph(g_run):
    run_fn()

CUPTI 레코드는 correlation id로 iteration에 귀속시켜, cast-copy 경로처럼 여러 커널이 나오는 경우도 전체 span을 잽니다.

li = bisect.bisect_left(launch_starts, t0)
ri = bisect.bisect_right(launch_starts, t1)
corr_ids = {sorted_launches[i].correlation_id for i in range(li, ri)}

정확도 게이트는 "fp32 reference 대비 커널 오차가 cast-copy 경로 자체 오차의 --tol-mult배 이하"로, 절대 tolerance보다 공정한 기준입니다. 다만 리뷰어는 이 스크립트를 flashinfer_benchmark.py harness에 통합하라고 요청했습니다.

왜 이게 좋은가

• 성능: tp-4 기준 H100/GB200/B200/B300 모든 행에서 기존 C++ 구현과 동등하거나 최대 약 5% 빠릅니다(B200 최악 +0.1%). 호스트 오버헤드는 shape별 plan 캐시, 컴파일 시에만 pointer 객체 생성, raw stream handle 사용으로 호출당 12 us.
• 정밀도: scale을 FP32 accumulator에 적용해 cast-copy의 이중 반올림을 피합니다.
• 유지보수: CUTLASS 포크가 사라져 업스트림 버전 업에 막히지 않습니다. 쓰이지 않던 tile(2sm256x256, 128x160)도 리뷰 후 제거됐습니다.

일반적 교훈

1. 가중치 전용 양자화는 dequantize를 커널에 fuse해야 이득이 실현된다. 사본을 만들면 메모리 절감이 사라집니다.
2. L2 재사용은 순회 방향 하나로 크게 바뀐다. 작은 operand를 빠르게 돌려 큰 operand를 L2에 머물게 하세요.
3. "컴파일러가 알아서 해줘서 맞는" 코드는 버그다. PTX 규약(wgmma.fence)을 명시적으로 지키세요.
4. JIT 캐시는 테스트를 거짓 통과시킬 수 있다. 정렬·컴파일 경로 테스트는 캐시를 끄고 돌려야 합니다.
5. co-residency를 가정하는 grid 동기화는 위험하다. bounded wait나 cooperative launch로 진행성을 보장하세요.

참고 자료

• PTX ISA (wgmma.fence, wgmma.mma_async) — Hopper mainloop의 fence 규약 근거
• torch.cuda.CUDAGraph — 벤치마크의 graph 캡처/replay
• torch.cuda.graph — L2 flush와 측정 op 캡처에 사용한 context manager
• torch.nn.functional.linear — cast-copy baseline 경로

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