ntransformer - Engine suy luận NVMe-to-GPU chạy Llama 3.1 70B trên một RTX 3090 duy nhất

• Engine suy luận LLM dựa trên C++/CUDA cho phép chạy mô hình Llama 70B trên RTX 3090 (24GB VRAM) thông qua streaming bộ nhớ GPU và I/O trực tiếp với NVMe
• Sử dụng cấu trúc cache thích ứng 3 tầng để tự động phân chia giữa VRAM, RAM cố định và NVMe/mmap, đạt tốc độ nhanh hơn tối đa 83 lần so với mmap
• Backend gpu-nvme-direct bỏ qua CPU hoàn toàn, truyền dữ liệu trực tiếp từ NVMe tới GPU để tận dụng tối đa băng thông PCIe
• Các tính năng Layer skip và self-speculative decoding giúp giảm tính toán không cần thiết và tăng tốc xử lý mà không làm giảm chất lượng
• Giúp vận hành hiệu quả các mô hình siêu lớn trên phần cứng tiêu dùng, mở ra khả năng mở rộng khả năng tiếp cận suy luận LLM hiệu năng cao

---

Tổng quan về NTransformer

• Engine suy luận LLM C++/CUDA hiệu suất cao chạy mô hình Llama 70B trên RTX 3090 (24GB VRAM)
  • Stream các layer mô hình qua bộ nhớ GPU và có thể tùy chọn dùng I/O trực tiếp với NVMe để bỏ qua CPU hoàn toàn
• Không có phụ thuộc bên ngoài ngoài CUDA Toolkit, không cần PyTorch hay cuBLAS
• Hỗ trợ định dạng mô hình GGUF, dùng được các kiểu lượng tử hóa Q4_0, Q8_0, Q4_K_M, Q5_K, Q6_K, F16, F32

Hiệu năng và cấu trúc cache

• Cache thích ứng 3 tầng (3-Tier Adaptive Caching)
  • Layer thường trú trong VRAM (0 I/O)
  • RAM cố định (chỉ dùng cho truyền H2D)
  • Dự phòng bằng NVMe/mmap
• Trong môi trường RTX 3090 + 48GB RAM, đạt tốc độ nhanh hơn 83 lần so với mmap
• Băng thông PCIe Gen3 x8 (khoảng 6.5 GB/s) là nút thắt cổ chai
• Lượng tử hóa Q4_K_M nạp thêm 10 layer vào VRAM (36 so với 26), giúp giảm lượng dữ liệu phải truyền
• Layer skip (dựa trên độ tương đồng cosine) bỏ qua 20 trong 80 layer, giảm thiểu tổn thất chất lượng

Tính năng chính

• Streaming SLEP: chồng lấp đọc NVMe, DMA qua PCIe và tính toán GPU bằng double buffering
• Backend gpu-nvme-direct: đọc trực tiếp dữ liệu NVMe vào bộ nhớ cố định mà GPU có thể truy cập
• Self-speculative decoding: dùng các layer thường trú trong VRAM làm mô hình nháp, không cần thêm mô hình phụ
• Tự động chọn đường đi dữ liệu: thường trú trong VRAM > H2D từ RAM cố định > mmap cố định > memcpy của CPU
• Hỗ trợ kiến trúc Llama: gồm RoPE, GQA, SwiGLU, RMSNorm và cache KV

Yêu cầu hệ thống

• Linux (Ubuntu, kernel 6.17+), CUDA Toolkit 13.1, gcc/g++ 14, CMake 3.24+
• GPU Compute Capability 8.0+ (đã thử nghiệm trên RTX 3090)
• Khi dùng I/O trực tiếp với NVMe, cần SSD NVMe ở khe PCIe riêng và thư viện gpu-nvme-direct

Streaming trực tiếp NVMe

• Khi mô hình không vừa trong VRAM, hệ thống dùng đường đi trực tiếp NVMe → GPU để loại bỏ CPU hoàn toàn
  • Luồng dữ liệu: NVMe SSD → DMA → bộ nhớ staging cố định → PCIe H2D → bộ đệm GPU → tính toán
• Gắn NVMe với VFIO để truy cập trực tiếp từ không gian người dùng
• Mỗi layer (khoảng 670MB với 70B Q6_K) được đọc bằng khoảng 670 lệnh NVMe trong khoảng 202ms
• Đọc NVMe, H2D DMA và tính toán GPU được xử lý song song bằng pipeline double buffering

Thiết lập hệ thống và cảnh báo rủi ro

• Script thiết lập tự động (setup_system.sh) cấu hình lần lượt GRUB, NVIDIA DKMS, CUDA header, VFIO, gắn NVMe
• Bao gồm các thao tác rủi ro cao như vô hiệu hóa IOMMU, vá module kernel, gắn NVMe với VFIO
• Cấu hình sai có thể gây không khởi động được, mất dữ liệu NVMe, hệ thống không ổn định
• Tuyệt đối không dùng ổ khởi động, cần một thiết bị NVMe riêng biệt
• Có cung cấp script sao lưu và khôi phục cho mọi thay đổi

Kiến trúc và cấu trúc mã nguồn

• Các thành phần chính trong thư mục src/
  • core/: tensor, cấp phát bộ nhớ, quản lý thiết bị GPU
  • cuda/: các kernel GEMV, RMSNorm, RoPE, SwiGLU, softmax
  • memory/: engine streaming SLEP dựa trên NVMe và mmap
  • model/: cấu trúc Transformer, bộ nạp GGUF, attention, FFN, normalization
  • inference/: tokenizer, sampler, engine
• scripts/: gồm các script thiết lập hệ thống, gắn NVMe và khôi phục

Lộ trình phát triển

• Giai đoạn 1: Llama 8B Q8_0, kernel CUDA tùy chỉnh, 48.9 tok/s (hoàn thành)
• Giai đoạn 2: streaming SLEP, chạy 70B trên một GPU, tăng tốc 33 lần (hoàn thành)
• Giai đoạn 3: hỗ trợ Q4_K_M/Q5_K, Layer skip, self-speculative decoding, cache KV F16 (hoàn thành)
• Giai đoạn 4: backend NVMe Direct, đọc NVMe do GPU điều phối ở mức 3.35 GB/s (hoàn thành)
• Giai đoạn 5: tối ưu suy luận và công khai C API (dự kiến)

Giấy phép

• Áp dụng giấy phép BSD-2-Clause

[GN⁺]

Ý kiến trên Hacker News

• Tôi nghĩ cách truyền trực tiếp từ NVMe sang GPU bỏ qua CPU thực sự rất thông minh
  Khi chạy các mô hình lớn ở máy cục bộ, nút thắt cổ chai luôn là hệ phân cấp bộ nhớ, còn cách này về cơ bản là dùng NVMe như VRAM mở rộng và truy cập trực tiếp qua DMA
  Tôi khá tò mò nếu so với cách tiếp cận bộ nhớ hợp nhất (unified memory) của Apple M series thì sẽ thế nào. M4 Max có thể nạp toàn bộ mô hình 70B vào bộ nhớ nhưng thông lượng lại thấp hơn 3090
  Sẽ rất thú vị nếu có benchmark so sánh hiệu năng gốc của 3090 dùng cách tiếp cận NVMe với M4 Max theo tiêu chí batch inference

  • Tôi có M3 nên định sẽ thử trên Metal

• Dùng GPUdirect thì có thể truyền DMA trực tiếp tới thiết bị lưu trữ
  Tôi tự hỏi sẽ thế nào nếu bộ nhớ m.2 thực chất là DRAM. Khi spill mô hình từ GPU thì đâu cần tính bền vững, như vậy có thể để lại RAM hệ thống cho CPU dùng

  • Dùng RAM disk có lẽ sẽ còn tốt hơn nhiều. Thật tiếc vì Intel Optane đã không trở thành tiêu chuẩn. Đó từng là công nghệ rất hợp với kiểu workflow này
  • Tôi cũng nghĩ vậy. Trước đây tôi từng có bài trình bày về dask-cudf tại Dask Summit, trong đó tăng tốc phân tích log bằng cấu trúc mảng SSD song song → GPUDirect Storage → PCIe → GPU A100. Giờ đem cấu trúc như vậy áp dụng cho LLM hoặc mô hình MoE có vẻ sẽ rất thú vị
  • Thực ra bộ nhớ m.2 dựa trên DRAM đã tồn tại dưới dạng CXL (Compute Express Link). Chỉ có điều giá RAM quá đắt, nên khó tận dụng băng thông 31GB/s trên mỗi đầu nối NVMe

• Tốc độ 0.2 token mỗi giây thì chậm cho chat nhưng vẫn đủ cho tác vụ batch/bất đồng bộ
  Tôi đang chạy một pipeline tạo nội dung tự động, trong đó nhiều lời gọi LLM được thực hiện đồng thời. Khâu tạo ảnh mới là nút thắt cổ chai nên toàn bộ công việc vốn cũng mất khoảng 20 phút
  Nếu có thể chạy mô hình 70B cục bộ thì sẽ tiết kiệm được chi phí token API, tức là giảm chi phí đáng kể

  • Xét về chi phí thì không hiệu quả. Với mức 0.5 tok/s thì mỗi giờ chỉ được 3600 token, trong khi hệ thống 3090 tiêu thụ 200~300W. Nếu chạy cùng lượng token đó qua llama 3.1 trên OpenRouter thì còn rẻ hơn tiền điện nhiều. Dù vậy, xét về suy luận bảo vệ quyền riêng tư thì vẫn có ý nghĩa
  • Cũng nên tính đến chi phí điện năng khi chạy 3090 ở mức 350W trong thời gian dài

• 0.2 tok/s thì ổn cho mục đích thử nghiệm, nhưng không đủ để dùng tương tác
  Trong đa số trường hợp, các mô hình 8B hoặc 13B được lượng tử hóa tốt sẽ cho cân bằng độ trễ-chất lượng tốt hơn

  • Tôi cũng chỉ muốn thử xem có khả thi hay không. Trước đây trên PS2 tôi từng đạt 3000 token mỗi giây với transformer cổ điển, nhờ kiến trúc gửi lệnh trực tiếp từ bộ nhớ sang GPU. PC thông thường phải đi qua CPU nên chậm hơn. GPU chuyên nghiệp giải quyết được vấn đề này nhưng quá đắt
  • Dù vậy, trong một số tình huống thì chất lượng của mô hình lớn vẫn quan trọng hơn
  • Chạy bằng tổ hợp CPU+GPU sẽ nhanh hơn. Tôi đạt khoảng 1.5 tok/s với cấu hình 7950X+3090
  • Mãi đến khi xem bảng hiệu năng tôi mới hiểu mục trên cùng là mô hình 8B. 5 giây/token thì quá chậm. Với hệ thống 5950X+128GB RAM của tôi, có lẽ kết hợp CPU với GPU 3060 còn nhanh hơn. Tôi cũng khó tin với nhận định rằng giới hạn tính toán trên 3090 là 2 giây/token

• Đây thực sự là một thử nghiệm rất thú vị. Đáng lẽ tôi cũng nên làm cái này sớm hơn
  Tôi muốn biết các con số throughput thực tế so với băng thông lý thuyết của PCIe là bao nhiêu. Không rõ đây là vấn đề độ trễ hay vấn đề băng thông

  • Trên thực tế là nút thắt băng thông. Main B450 của tôi chỉ hỗ trợ PCIe3 x8 nên GPU bị giới hạn. Nếu nâng cấp lên X570 thì tốc độ có lẽ sẽ tăng 2~3 lần

• Một màn hack rất ngầu, nhưng 0.5 tok/s với mô hình 70B vẫn chậm nếu so với việc cùng chiếc card đó chạy mô hình 7B ở mức 30+ tok/s
  Theo nghiên cứu của NVIDIA, các mô hình dưới 10B cũng có thể xử lý 40~70% tác vụ agent, và khoảng cách chất lượng cũng đang thu hẹp rất nhanh

• Lĩnh vực này còn rất đáng để thử nghiệm trong tương lai
  Về dài hạn, điều cốt lõi có lẽ sẽ là tối ưu mô hình, tức nghiên cứu tìm ra những phần của mô hình có thể lược bỏ mà không ảnh hưởng đến hiệu năng. Rốt cuộc mô hình cũng là một dạng nén mất dữ liệu (lossy compression). Hướng đi này cũng sẽ giúp dân chủ hóa AI

  • Phép so sánh với nén khá thú vị. Fine-tuning cũng có thể nhìn theo cách tương tự. Ví dụ, nếu tinh chỉnh một mô hình 3B cho một tác vụ cụ thể thì sẽ không còn cần đến tính đa dụng của mô hình 70B nữa

• Dự án này thật sự rất ấn tượng. Tôi tò mò muốn biết cần nền tảng kiến thức hệ thống/phần cứng như thế nào mới có thể nghĩ ra ý tưởng kiểu này
  Tôi làm việc trong môi trường mà phần cứng đã bị trừu tượng hóa rất nhiều, nên rất khó nghĩ ra cách tiếp cận như vậy. Có lẽ không chỉ cần sáng tạo mà còn phải có hiểu biết ở cấp độ hệ thống

  • Đây chính là thử nghiệm đó: dự án ps2-llm
    Khi cố chạy LLM trên PS2, tôi đụng phải giới hạn RAM 32MB và VRAM 4MB nên đã nghĩ ra cách streaming layer. PS2 có thể để VRAM xử lý trực tiếp địa chỉ 32-bit nên rất nhanh, và tôi muốn tái hiện điều đó trên PC
  • Tôi cũng nghĩ tương tự. Console đời cũ có CPU chậm nên truyền DMA là yếu tố cốt lõi. Có lẽ chính trải nghiệm đó đã dẫn tới ý tưởng này. Thẻ nhớ thông minh trên PS2 cũng có chức năng DMA khá phức tạp

• Tôi cũng đang thử một thứ tương tự. Tôi đang làm thí nghiệm chạy mô hình 1T với chưa tới một nửa VRAM
  Có vẻ như nếu sửa lớp định tuyến của SGLang thì có thể triển khai expert swap dựa trên dự đoán JIT từ NVMe Gen5 sang bộ nhớ GPU. Tôi đang dùng các primitive của NVIDIA Dynamo và NIXL
  Không biết đã có ai thử việc này chưa

  • Tôi cũng rất muốn xem. Tôi đang cân nhắc mua thêm một chiếc 3090 nữa và nâng cấp bo mạch chủ để giải quyết nút thắt PCIe3. Có vẻ sẽ chạy được glm 4.7~5 với q4_k_m

• Dự án rất hay. Tôi muốn biết chi tiết hơn về quy trình vá DKMS trên GPU phổ thông. Tôi cũng muốn thử

  • Tôi đã cập nhật tài liệu để thêm quy trình vá và thông tin rủi ro
  • Cũng đã có những trường hợp sửa driver mã nguồn mở của NVIDIA để mở khóa các tính năng doanh nghiệp như giao tiếp P2P giữa GPU hoặc chia tách vGPU
    Liên kết liên quan: RTX4090 P2P Unlock, vGPU Unlock