61 · Compute-bound so với memory-bound và mô hình roofline

Nguồn: fanout.sh / inference-eng / 12-07 · ~6 phút

Ghi chú

Mọi kernel đều làm một ít phép tính và chuyển một ít byte, và mất thời gian xấp xỉ bằng cái lớn hơn trong hai. So arithmetic intensity của kernel (FLOP mỗi byte) với điểm gãy (ridge point) của chip (FLOPS đỉnh ÷ băng thông, ≈ 295 trên H100 ở BF16): dưới mức đó bạn memory-bound và thêm TFLOPS chẳng thay đổi gì; trên mức đó bạn compute-bound và thêm băng thông chẳng thay đổi gì. Roofline vẽ đúng điều này. GEMM của decode có intensity ≈ kích thước batch, nên chúng sống trên sườn dốc và batching gần như miễn phí; prefill sống trên mái.

1. Quy tắc max

Kernel nhân hai trên H100.

$$T_{\mathrm{math}}\ge \frac{\mathrm{FLOPs}}{\mathrm{peak}\mathrm{FLOPS}}{T}_{\mathrm{mem}}\ge \frac{\mathrm{bytes}}{\mathrm{bandwidth}}T\approx max(T_{\mathrm{math}},T_{\mathrm{mem}})$$

• Một kernel tốt chồng lấp hai việc, nên tổng xấp xỉ bằng cái lớn hơn.
• Phép tính lớn hơn → compute-bound. Byte lớn hơn → memory-bound.
• y = x * 2: phép tính ≈ 0,25 µs, byte ≈ 20 µs. Memory-bound áp đảo.

2. Hai tỷ số thay cho hai thời gian

Intensity của kernel so với điểm gãy của chip.

$$I=\frac{\mathrm{FLOPs}}{\mathrm{bytes}\mathrm{moved}}{I}_{\mathrm{ridge}}=\frac{989\mathrm{TFLOPS}}{3.35\mathrm{TB}/\mathrm{s}}\approx 295\mathrm{FLOP}/\mathrm{byte}$$

• Arithmetic intensity: một kernel làm bao nhiêu phép tính cho mỗi byte nó chạm tới. Nhân hai làm 1 FLOP mỗi 4 byte (đọc 2, ghi 2 ở BF16): I = 0,25.
• Tỷ số của chip: để giữ đơn vị tính toán luôn bận, mỗi byte mang vào phải trả cho ≈ 295 phép tính.
• Dưới điểm gãy: memory-bound. Trên: compute-bound.

3. Roofline

Cả hai trục theo thang log.

• Trục ngang: intensity. Trục dọc: FLOPS đạt được.
• Trần memory-bound: băng thông × intensity, một đường xiên.
• Trần compute-bound: đỉnh của chip, một mái phẳng.
• Trần của bạn là cái thấp hơn trong hai; chúng gặp nhau ở điểm gãy (295).
• Các op thật rơi vào đâu:
  • nhân từng phần tử (I = 0,25): tít dưới sườn dốc, bị chặn ở ≤ 0,84 TFLOPS;
  • softmax, layer norm, attention một token: vài FLOP mỗi byte, cũng nằm trên sườn dốc;
  • matmul vuông 4096 × 4096: I > 1.000 (◆ 2N³ FLOP ÷ 3 × 2N² byte ≈ 1.365), nằm trên mái.

4. Vì sao bức tranh này điều khiển inference

Panel phải là mô hình lý tưởng hoá của tôi cho thời gian bước của một GEMM decode.

$$\mathrm{FLOPs}=2BDF\mathrm{bytes}\approx 2DF\Rightarrow I\approx B$$

• Trong decode, các kernel lớn nhất nhân B vector token với một ma trận trọng số D × F khổng lồ.
• FLOP tăng theo batch, nhưng byte gần như toàn là trọng số, đọc một lần bất kể batch. Vì vậy intensity ≈ kích thước batch.
• B = 1: I ≈ 1, bị chặn ở 3,35 TFLOPS, khoảng ⅓ của 1% đỉnh.
• Tăng batch: điểm trượt lên dọc sườn dốc; mỗi bước cho nhiều token hơn với thời gian gần như không đổi.
• B = 64: vẫn memory-bound; 64 token thay vì 1 với thời gian gần như như cũ.
• Vài trăm (≈ 295): điểm chạm tới điểm gãy và batching không còn miễn phí.
• Prefill là thái cực bên kia: hàng nghìn token prompt là một batch khổng lồ, nằm trên mái.

5. Dùng roofline

Hai chẩn đoán khác nhau, hai nước đi khác nhau.

• Tìm điểm của bạn. Biểu đồ roofline của Nsight Compute vẽ sẵn từng kernel cho bạn.
• Bên trái điểm gãy: phép tính của một chip nhanh hơn không giúp được gì. Dịch sang phải:
  • batch nhiều việc lại với nhau;
  • gộp (fuse) kernel để byte không bị nạp lại;
  • lưu trọng số bằng ít bit hơn: FP8 giảm một nửa số byte và nhân đôi intensity.
• Nằm thấp hẳn dưới đường: phần cứng không phải giới hạn, kernel mới là giới hạn. Đó là việc của code tốt hơn.

6. Hai lời cảnh báo

Mái theo datasheet là cận trên.

• Mái thật thấp hơn mái trong datasheet: H100 thường đạt ≈ 80% FLOPS đỉnh. ◆ Ở 80% (≈ 791 TFLOPS), điểm gãy dịch vào ≈ 236.
• Mỗi độ chính xác có mái riêng. Ở FP8 đỉnh tăng gấp đôi lên 1.979 TFLOPS, nên điểm gãy dịch ra ≈ 590.
• ◆ Trọng số FP8 nhân đôi intensity của một GEMM decode (I ≈ 2B), nên trên mái FP8, điểm gãy vẫn đạt được ở batch ≈ 295. Với trọng số FP8 nhưng phép tính BF16, nó đạt ở ≈ 148.

Quan trọng

Biết một kernel nằm phía nào của điểm gãy trước khi tối ưu. Bên trái thì mua băng thông hoặc dịch sang phải; bên phải thì mua FLOPS; nằm thấp hẳn dưới đường thì sửa code.

7. Tóm tắt