DeepGEMM: kernel FP8 GEMM sạch và hiệu quả nhờ scaling chi tiết

 (github.com/deepseek-ai)
2 điểm bởi GN⁺ 2025-02-27 | 1 bình luận | Chia sẻ qua WhatsApp
  • DeepGEMM là thư viện kernel tensor core hiệu năng cao, gom các primitive tính toán cốt lõi của LLM hiện đại như GEMM, fused MoE, MQA scoring, HyperConnection... vào một codebase CUDA duy nhất
  • Tất cả kernel được biên dịch lúc chạy bằng mô-đun JIT gọn nhẹ; không cần biên dịch CUDA khi cài đặt và yêu cầu C++20, CUDA Toolkit, PyTorch, CUTLASS 4.0 trở lên
  • Dù tận dụng một số khái niệm của CUTLASS và CuTe, thư viện không phụ thuộc nặng vào template hay cấu trúc đại số, và được thiết kế để việc học tối ưu hóa kernel NVIDIA GPU dễ tiếp cận hơn thông qua số lượng hàm kernel cốt lõi giới hạn
  • Phạm vi hỗ trợ bao gồm FP8, FP4, BF16 GEMM, grouped GEMM, kernel MQA logits cho DeepSeek v3.2, đến Mega MoE chồng lấp giao tiếp và tính toán; các ràng buộc về layout bộ nhớ khác nhau giữa SM90 và SM100
  • Dù thiết kế gọn nhẹ, thư viện hướng tới hiệu năng ngang bằng hoặc cao hơn các thư viện được tinh chỉnh chuyên sâu trên nhiều shape ma trận, và có bản cập nhật đạt tối đa 1550 TFLOPS trên H800

Mục tiêu và thiết kế của DeepGEMM

  • DeepGEMM là thư viện kernel tensor core tích hợp các primitive tính toán chính dùng trong các mô hình ngôn ngữ lớn hiện đại vào một codebase CUDA duy nhất
    • GEMM: FP8, FP4, BF16
    • fused MoE chồng lấp giao tiếp: Mega MoE
    • MQA scoring cho lightning indexer
    • HyperConnection(HC)
  • Tất cả kernel được biên dịch lúc chạy dưới dạng mô-đun Just-In-Time(JIT) gọn nhẹ
    • Không cần biên dịch CUDA trong quá trình cài đặt
  • Tận dụng một số khái niệm của CUTLASS và CuTe
    • Tuy nhiên không phụ thuộc nhiều vào template nặng hay cấu trúc đại số
    • Giữ codebase đơn giản bằng cách giới hạn số lượng hàm kernel cốt lõi
  • Dù có thiết kế gọn nhẹ, dự án cho biết có thể đạt hiệu năng ngang bằng hoặc cao hơn các thư viện được tinh chỉnh chuyên sâu trên nhiều shape ma trận

Các cập nhật chính

  • Bản cập nhật ngày 16/4/2026 bao gồm Mega MoE, FP8xFP4 GEMM, FP4 Indexer, PDL, biên dịch JIT nhanh hơn...
    • Chi tiết tại #304
    • Benchmark Mega MoE tại #316
  • Bản cập nhật ngày 28/9/2025 bổ sung kernel scoring weighted ReLU MQA logits cho DeepSeek v3.2 lightning indexer
    • Chi tiết tại #200
  • Bản cập nhật ngày 20/7/2025 hỗ trợ cả SM90 và SM100, đồng thời refactor toàn bộ sang mô-đun JIT CPP có overhead CPU thấp
    • NVRTC và tối ưu hóa SASS sau biên dịch bị vô hiệu hóa
    • NVRTC được đánh dấu là sẽ được hỗ trợ sau
    • Do NVCC 12.9 tự động thực hiện FFMA interleaving, tối ưu hóa sau biên dịch không còn được hỗ trợ
    • Chi tiết tại #112
  • Bản cập nhật ngày 14/5/2025 bổ sung kernel weight gradient cho dense và MoE backward
    • Chi tiết tại #95
  • Bản cập nhật ngày 7/5/2025 cung cấp tốc độ biên dịch nhanh hơn tới 10 lần nhờ hỗ trợ NVRTC
    • Có thể bật bằng DG_JIT_USE_NVRTC=1
    • Trong một số trường hợp có thể bị giảm hiệu năng
    • Chi tiết tại #94
  • Bản cập nhật ngày 18/4/2025 đạt tối đa 1550 TFLOPS trên H800

Yêu cầu và quy trình cài đặt

  • Môi trường chạy yêu cầu GPU kiến trúc NVIDIA SM90 hoặc SM100
  • Yêu cầu phần mềm như sau
    • Python 3.8 trở lên
    • Trình biên dịch hỗ trợ C++20
    • CUDA Toolkit
      • SM90: CUDA 12.3 trở lên
      • Khuyến nghị mạnh CUDA 12.9 trở lên để đạt hiệu năng cao nhất
      • SM100: CUDA 12.9 trở lên
    • PyTorch 2.1 trở lên
    • CUTLASS 4.0 trở lên
    • Thư viện {fmt}
  • Trong môi trường phát triển, clone repository kèm submodule rồi chạy develop.sh để liên kết các include cần thiết và build mô-đun CPP JIT
  • Cài đặt bằng cách chạy install.sh, sau đó import deep_gemm trong dự án Python

Giao diện GEMM và ràng buộc layout

  • Quy ước đặt tên kernel GEMM của DeepGEMM là D = C + A @ B
  • Layout shape đầu vào dựa trên NT
    • fp8_gemm_nt thực hiện D = C + A @ B.T
  • Triển khai SM90 chỉ hỗ trợ layout bộ nhớ NT
    • Tương ứng với tổ hợp row-major, col-major
  • Triển khai SM100 hỗ trợ toàn bộ các layout bộ nhớ NT, TN, NN, TT
  • Trên cả hai kiến trúc, LHS scaling factor phải được căn chỉnh TMA và ở layout đã transpose
    • SM90 yêu cầu scaling factor ở định dạng FP32
    • SM100 yêu cầu định dạng packed UE8M0, đóng gói 4 UE8M0 vào một torch.int
  • Các tác vụ như transpose đầu vào hoặc cast sang FP8 cần được người dùng xử lý riêng
    • Thư viện cung cấp các hàm tiện ích PyTorch đơn giản nhưng có thể chậm
    • Trọng tâm chính là tối ưu hóa kernel GEMM

Dense và Grouped GEMM

  • FP8 GEMM non-grouped cơ bản dùng các hàm fp8_gemm_{nt, nn, tn, tt}
  • Grouped GEMM với contiguous layout chỉ group theo trục M, khác với grouped GEMM truyền thống của CUTLASS
    • N và K phải cố định
    • Thiết kế này phù hợp với trường hợp các expert trong mô hình MoE dùng chung shape
  • Trong training forward pass hoặc inference prefilling, số token được xử lý bởi mỗi expert có thể khác nhau
    • Dạng nối các token này thành một tensor duy nhất được gọi là contiguous layout
    • Mỗi segment expert phải được căn chỉnh theo GEMM M block size
    • Có thể kiểm tra tiêu chí căn chỉnh bằng get_mk_alignment_for_contiguous_layout()
  • Cũng cung cấp API grouped theo trục K cho MoE weight backward
    • M và N phải cố định
    • Hàm liên quan là k_grouped_fp8_gemm_tn_contiguous
  • Ở giai đoạn inference decoding, khi CUDA graph được bật và CPU không thể biết số token theo từng expert, thư viện hỗ trợ masked grouped GEMM
    • Nếu cung cấp mask tensor, kernel chỉ tính vùng hợp lệ
    • Hàm là m_grouped_fp8_gemm_nt_masked
    • Có ví dụ dùng output của kernel low-latency từ DeepEP làm input

Kernel MQA cho DeepSeek v3.2 Indexer

  • Họ kernel V3.2 MQA cung cấp phiên bản non-paged và paged
    • non-paged dùng cho prefilling
    • paged dùng cho decoding
  • fp8_mqa_logits nhận 6 input
    • q: tensor E4M3, shape [seq_len, num_heads, head_dim]
    • kv: tensor E4M3 và float scaling factor
      • shape tensor là [seq_len_kv, head_dim]
      • shape scaling factor là [seq_len_kv]
    • weights: tensor float, shape [seq_len, num_heads]
    • cu_seq_len_k_start, cu_seq_len_k_end: tensor int, shape [seq_len]
    • clean_logits: có dọn các logits chưa được lấp đầy thành -inf hay không
  • Shape tensor output là [seq_len, seq_len_kv] và biểu diễn token-to-token logits
  • Mỗi q token i duyệt qua các kv token j từ cu_seq_len_k_start[i] đến trước cu_seq_len_k_end[i]
    • Nhân kv_j với scaling factor
    • Tính giá trị theo từng head bằng q[i, :, :] @ kv_j
    • Sau khi áp dụng ReLU, nhân với weights[i, :] rồi cộng lại để tạo scalar logit
  • Hàm phiên bản paged là fp8_paged_mqa_logits

Mega MoE

  • Mega MoE fuse nhiều bước MoE vào một mega-kernel
    • EP dispatch
    • linear 1, FP8xFP4
    • SwiGLU
    • linear 2, FP8xFP4
    • EP combine
  • Mega MoE overlap giao tiếp NVLink với tính toán tensor core
  • Việc chạy yêu cầu multi-process launch dùng symmetric memory
  • Quy trình sử dụng như sau
    • Cấp phát symmetric memory buffer bằng deep_gemm.get_symm_buffer_for_mega_moe
      • Yêu cầu PyTorch 2.9 trở lên
    • Chuyển đổi weight gồm FP4 và UE8M0 SF sang layout yêu cầu bằng deep_gemm.transform_weights_for_mega_moe
    • Trước khi gọi, sao chép input, scaling factor, top-k index, top-k weight vào buffer
    • Chạy kernel fused mega MoE bằng deep_gemm.fp8_fp4_mega_moe
  • Thiết lập multi-process đầy đủ và ví dụ benchmark nằm trong tests/test_mega_moe.py

Tiện ích và biến môi trường

  • Các hàm tiện ích chính kiểm soát tài nguyên chạy, căn chỉnh, biên dịch JIT và chuyển đổi scaling factor
    • deep_gemm.set_num_sms / get_num_sms: đặt và truy vấn số SM tối đa sẽ dùng
    • deep_gemm.set_tc_util / get_tc_util: đặt và truy vấn tỷ lệ tensor core utilization xấp xỉ
    • deep_gemm.set_pdl / get_pdl: bật và tắt Programmatic Dependent Launch(PDL)
    • deep_gemm.set_mk_alignment_for_contiguous_layout / get_mk_alignment_for_contiguous_layout: đặt và truy vấn group-level M/K alignment của contiguous layout
    • deep_gemm.transform_sf_into_required_layout: chuyển scaling factor sang layout yêu cầu
    • deep_gemm.get_tma_aligned_size: truy vấn kích thước TMA alignment cần thiết
  • Các biến môi trường liên quan đến JIT kiểm soát output debug, vị trí cache, lựa chọn compiler và tùy chọn profiling
    • DG_JIT_DEBUG: xuất thông tin debug JIT
    • DG_PRINT_CONFIGS: xuất config được chọn theo từng shape
    • DG_JIT_CACHE_DIR: thư mục cache kernel đã biên dịch, mặc định là $HOME/.deep_gemm
    • DG_JIT_USE_NVRTC: dùng NVRTC thay cho NVCC, cho phép biên dịch nhanh nhưng trong một số trường hợp hiệu năng có thể thấp hơn
    • DG_JIT_NVCC_COMPILER: đường dẫn compiler NVCC
    • DG_JIT_CPP_STANDARD: phiên bản chuẩn C++, mặc định là 20
  • Cũng cung cấp các biến môi trường debug và profiling
    • DG_JIT_DUMP_ASM, DG_JIT_DUMP_PTX, DG_JIT_DUMP_SASS: dump output PTX và SASS
    • DG_JIT_WITH_LINEINFO: bao gồm thông tin source line cho công cụ profiling
    • DG_COMM_KERNEL_DEBUG: khởi tạo symmetric buffer về 0 trước khi gọi Mega MoE
    • DG_USE_NVIDIA_TOOLS: bỏ qua profiling nội bộ khi chạy công cụ NVIDIA bên ngoài
  • Tùy chọn build kiểm soát cách cài đặt và tải kernel
    • DG_SKIP_CUDA_BUILD: bỏ qua CUDA extension build trong khi cài đặt
    • DG_FORCE_BUILD: buộc build cục bộ thay vì tải pre-built wheel
    • DG_JIT_USE_RUNTIME_API: dùng CUDA Runtime API để tải kernel, yêu cầu CUDA runtime 12.8 trở lên

Giấy phép và trích dẫn

  • Repository DeepGEMM được phát hành theo MIT License
  • Dự án cho biết được truyền cảm hứng từ CUTLASS
  • Tiêu đề mục trích dẫn là DeepGEMM: clean and efficient BLAS kernel library on GPU

Bài viết liên quan

1 bình luận

 
GN⁺ 2025-02-27
Ý kiến trên Hacker News

  • Interleaving FFMA SASS trông thật sự đáng kinh ngạc
    Có vẻ như khi thấy hiệu năng kernel CUTLASS FP8 được cải thiện giữa NVCC 12.2 và 12.3, họ đã so sánh SASS đã biên dịch và phát hiện một bit trong nhiều lệnh FADD bị đảo theo mẫu interleaving; rồi tham khảo triển khai CUDA assembler mã nguồn mở để xác định bit đó là yield bit, cho phép warp hiện tại nhường quyền để warp khác chạy
    Ấn tượng ở chỗ họ đã tận dụng điều này để tạo một script sửa các lệnh FFMA trong binary đã biên dịch; không chỉ yield bit mà cả reuse bit cũng được đảo cùng lúc, vì khi warp nhường quyền thì không thể tái sử dụng thanh ghi, nhờ đó các lệnh MMA và các lệnh FFMA thăng cấp trong FP8 GEMM với scaling tinh vi chồng lấp tốt hơn, tăng hiệu năng hơn 10% trong một số trường hợp

    • Theo những gì tôi đọc được ở nơi khác, trong tối ưu hóa phép toán ma trận nhạy cảm về hiệu năng, cách làm này khá điển hình
      Chỉ là với vấn đề cụ thể này, có vẻ các công ty AI khác chưa thấy cần thiết nên chưa áp dụng; cuối cùng nhiều khả năng mọi người cũng sẽ đi tới những điểm tương tự
    • Scott Gray đã tìm ra đúng những thứ như thế này, và còn hơn thế nữa, trên Maxwell vào năm 2015; sau đó cũng đã có nhiều người bàn đến khá nhiều

  • Những trường hợp như thế này cho thấy các trình biên dịch hiện nay còn cách rất xa việc khai thác hiệu năng phần cứng chỉ bằng mã cấp cao
    Tôi tự hỏi cần những gì để các kỹ thuật biên dịch truyền thống hoặc các agent tối ưu hóa dựa trên AI tạo ra được kết quả như vậy

    • Có lẽ sẽ cần một lượng thử-sai khổng lồ trong vòng lặp phản hồi học tăng cường

  • Con số tăng tốc được báo cáo là so với baseline tự xây dựng dựa trên CUTLASS
    Tôi tò mò không biết đã có ai so sánh hiệu năng trực tiếp với cuBLAS chưa
    Các kết quả CUTLASS GEMM tôi từng thấy cho tới nay thường trong khoảng 10% so với cuBLAS; nếu mức cải thiện 2x~2,5x như bài viết nói vẫn giữ được thì sẽ thật sự ấn tượng

    • Thường thì tôi tránh FP8 và chuộng I8 hơn, nhưng câu hỏi này khiến tôi tò mò cuBLAS chạy tốt đến mức nào
      Trước hết, cuBLAS cần API mở rộng cuBLASLt để xử lý các tác vụ mixed precision như FP8
      Ngoài ra, nó không hỗ trợ các tổ hợp kiểu tưởng như phù hợp như E5M2 x E5M2 trong A x B, nhưng có hỗ trợ E5M2 x E4M3; và các ràng buộc vẫn tiếp diễn, chẳng hạn trên Ampere, Hopper, Blackwell, ma trận A luôn phải ở layout đã chuyển vị
      Tôi đã tích hợp benchmark FP8 cuBLASLt vào repo "Less Slow C++" của mình <https://github.com/ashvardanian/less_slow.cpp>, và thêm nó vào danh sách benchmark cuBLAS hiện có cùng các benchmark CUDA/PTX tự viết
      Tôi đang chạy trên GPU H200, vốn đáng lẽ có hiệu năng tương đương H100; với đầu vào vuông, throughput đạt đỉnh khoảng 1,35 Peta-ops
      256 là 2,68T/s, 512 là 20,49T/s, 1024 là 144,23T/s, 2048 là 665,68T/s, 4096 là 1,26P/s, 8192 là 1,34P/s, 16384 là khoảng 1,23P/s; con số này bằng khoảng 67% mức NVIDIA quảng bá cho dense GEMM <https://resources.nvidia.com/en-us-data-center-overview-mc/e...>
    • Tôi nghe nói có thể đạt hiệu năng tốt hơn cuBLAS bằng CUTLASS
      Tôi đã nghĩ baseline là chọn bên tốt hơn giữa cuBLAS và CUTLASS

  • Mã nguồn mở kiểu này thật sự cho thấy rất rõ mục tiêu đạt hiệu quả của ngành
    Tuy nhiên, lợi ích từ phần mềm này có lẽ sẽ đến với các công ty lớn đang serving mô hình ở quy mô lớn — tức những đối thủ tiềm năng của DeepSeek — nhiều hơn là với cộng đồng mã nguồn mở phổ thông muốn học hỏi, thử nghiệm hoặc serving mô hình trên phần cứng tiêu dùng

    • Khi hiệu quả tăng lên, cuối cùng nó có thể dẫn tới phần cứng rẻ hơn cho tất cả mọi người, kể cả chính DeepSeek

  • Tôi không chắc việc tối ưu theo hướng độ chính xác ngày càng thấp có tốt về dài hạn hay không
    Điều này có nghĩa là mô hình thực ra khá thưa, và dù hiện tại có thể đúng như vậy, tôi cho rằng khả năng cao là do cách huấn luyện đang trộn lẫn một số ý tưởng không hay, hơn là vì về bản chất nó phải thưa như thế

    • Trong khi tính thưa miễn phí còn hiệu quả thì cứ tận hưởng
      Làm sao để chỉ ở độ chính xác cao hơn mới huấn luyện được các mô hình rất tốt là một vấn đề nghiên cứu, còn huấn luyện và suy luận độ chính xác thấp là vấn đề kỹ thuật
      Chúng ta đã làm việc này từ thời CNN, ít nhất là 9 năm trước, và tôi nghĩ vẫn còn vài năm nữa
    • Vì hàm kích hoạt loại bỏ khá nhiều dải động của số dấu phẩy động, khá rõ là việc dành một dải rộng cho các vùng kích hoạt vốn đã bão hòa có thể không hữu ích

  • Điều này có thể trở nên vô nghĩa vì MXFP, tức hỗ trợ microscaling native của Blackwell
    Trên Hopper, về cơ bản họ triển khai thủ công ở mức hạt thô hơn, nhưng dùng hệ số scaling FP32

    • Đúng vậy
      Những bản trình diễn công khai chất lượng cao như thế này cho thấy rõ hào lũy của $NVDA nằm ở đâu
      GPU đa dụng rất linh hoạt, nên có thể lập trình để thực hiện nhiều tác vụ hợp lý mà nhà cung cấp phần cứng ban đầu chưa nghĩ tới
      Nhưng nếu dự đoán tương lai sẽ ngày càng hội tụ về hỗ trợ phần cứng chuyên dụng hơn, khiến dư địa tối ưu phần mềm kiểu này biến mất, thì cái gọi là hào lũy CUDA sẽ sụp đổ
      Để tiếp tục ở lại cuộc chơi này, NVIDIA đang tự phá bỏ hào lũy của mình :p

  • Chà, là giấy phép MIT
    Mong các công ty lớn sẽ chấp nhận kiểu hợp tác mã nguồn mở như thế này

  • Tôi vẫn luôn thắc mắc vì sao lại tồn tại các lệnh không được tài liệu hóa
    Dù không hoàn toàn ổn định, tôi vẫn nghĩ có lẽ nên cung cấp chúng cho người dùng
    Những nội dung như vậy chắc hẳn được tài liệu hóa nội bộ, nhưng tôi không hiểu vì sao lại không công khai
    Bảo mật dựa vào sự mơ hồ thì không hiệu quả, còn đối thủ cạnh tranh thì kiểu gì cũng reverse engineer hết

    • Có lẽ cũng tương tự lý do vì sao những thứ chúng ta làm ra cũng có các phần không được tài liệu hóa
      Có thể là do thiếu thời gian, hoặc vì không muốn ám chỉ rằng mình sẽ hỗ trợ các tính năng chưa ổn định hay còn mang tính thử nghiệm
      Nếu thiệt hại chỉ dừng ở mức đội ngồi bên cạnh thì việc thay đổi cũng dễ hơn nhiều
    • Ngay từ tiền đề “những nội dung này chắc được tài liệu hóa nội bộ” có thể đã không đúng
      Có khả năng cao chúng chỉ nằm trong những nơi như tài liệu thiết kế kiến trúc hoặc đặc tả, và những tài liệu như vậy thì đương nhiên họ sẽ không muốn chia sẻ

  • Thành thật mà nói, nội dung này vượt quá phạm vi sử dụng và hiểu biết của tôi
    Dù vậy, tôi thật sự biết ơn và thấy mới mẻ khi những phát hiện và cải tiến như thế này được chia sẻ để mọi người cùng hưởng lợi

    • FFMA là viết tắt của Fused Floating-point Multiply-Add, một lệnh GPU cơ bản thực hiện D = A*B + C trong một lần
      Nó rất quan trọng trong phép nhân ma trận và workload học sâu
      Trong SASS của NVIDIA, lệnh FFMA được mã hóa dưới dạng lệnh 64 bit hoặc 128 bit, và có nhiều bit điều khiển xác định hành vi chính xác
      Khi bit yield được đặt, nó báo cho warp scheduler rằng sau lệnh này warp hiện tại có thể nhường quyền thực thi, và phần cứng có thể chạy warp khác để che giấu độ trễ
      GPU đạt thông lượng cao nhờ tính song song quy mô lớn, và khi một warp bị dừng vì chờ bộ nhớ chẳng hạn, warp khác có thể tiếp tục tiến hành
      Bit reuse cho biết liệu thanh ghi nguồn có thể được tái sử dụng ngay trong phép toán kế tiếp hay không, và nếu bit yield được đặt thì nó bắt buộc phải tắt
      Lý do là khi một warp nhường quyền, warp được thực thi tiếp theo có thể không phải là chính nó, và warp khác có thể thay đổi trạng thái file thanh ghi, nên phần cứng không thể bảo đảm rằng giá trị thanh ghi được giữ nguyên vượt qua điểm yield
      Nếu đặt bit yield cho các lệnh FFMA theo mẫu xen kẽ, trình biên dịch sẽ tạo ra các điểm lập lịch tường minh để warp khác có thể tiến hành, và để giữ tính đúng đắn cũng phải xóa bit reuse của lệnh đó
      Thay đổi này đặc biệt hữu ích cho việc chồng lấp lệnh MMA, vốn là lõi của phép nhân ma trận, với lệnh FFMA nâng cấp dùng để chuyển đổi FP8 nhằm tích lũy ở độ chính xác cao hơn
      FP8 GEMM thường cần quá trình chuyển đổi sang độ chính xác cao hơn để tích lũy rồi chuyển ngược lại, tạo ra các FFMA bổ sung; điều này làm giảm yêu cầu băng thông bộ nhớ nhưng tạo nên một mẫu tính toán phức tạp xen lẫn các phép nâng cấp/hạ cấp
      “Scaling tinh vi” có vẻ chỉ việc quản lý độ chính xác một cách cẩn thận ở nhiều điểm trong quá trình tính toán
      Việc thao tác bit yield giúp các phép toán tính toán và chuyển đổi định dạng được interleave tốt hơn, khiến các đơn vị thực thi của GPU được dùng hiệu quả hơn; nếu không có tối ưu này, warp scheduler có thể không tìm được cơ hội chuyển đổi tự nhiên, dẫn đến tài nguyên tính toán được tận dụng kém hơn