HBF dành cho trọng số mô hình AI, HBM dành cho KV Cache

Hãy chuẩn bị cho một bước chuyển lớn trong hạ tầng bộ nhớ AI. Nút thắt về bộ nhớ, đặc biệt là HBM, có thể sẽ được giải tỏa đáng kể nhờ sáng kiến phát triển High-Bandwidth Flash (HBF) của SanDisk cùng SK hynix. Công nghệ này được kỳ vọng sẽ phân tách các tác vụ vốn trước đây đều do HBM đảm nhiệm.

Mặc dù HBF có thể cung cấp băng thông đọc cao hơn nhiều lần so với chuẩn JEDEC HBM4, công nghệ này vẫn gặp hạn chế về tốc độ ghi và độ bền ghi.

HBM và bài toán "bức tường bộ nhớ"

Một mô hình AI có 100 tỷ tham số và hoạt động ở độ chính xác FP16 (2 byte cho mỗi tham số) sẽ cần khoảng 200 GB HBM chỉ để lưu trữ trọng số (weights) của mô hình.

Đối với những ai chưa quen với khái niệm này, trọng số là thành phần quyết định mức độ quan trọng mà mô hình gán cho từng khái niệm hoặc chuỗi từ. Nói cách khác, đây là toàn bộ kiến thức mà mô hình đã học được. Thông thường, mô hình càng lớn thì càng cần nhiều trọng số để lưu trữ.

Tuy nhiên, đó mới chỉ là một phần của vấn đề. Hãy tưởng tượng bạn đang viết một câu chuyện nhưng lại có trí nhớ ngắn hạn rất kém. Mỗi khi viết thêm một từ mới, bạn đều phải đọc lại toàn bộ những gì đã viết trước đó để nhớ nội dung. Văn bản càng dài thì việc này càng mất thời gian.

KV Cache hoạt động giống như việc ghi chú ra một tờ giấy riêng để luôn nhớ được những gì đã xử lý trước đó. Nhờ vậy, quá trình suy luận của AI được tăng tốc đáng kể. Tuy nhiên, khi cửa sổ ngữ cảnh (context window) ngày càng lớn, dung lượng KV Cache cũng tăng theo. Do AI phải truy cập KV Cache liên tục với tốc độ rất cao nên dữ liệu này thường được lưu trong HBM.

Đây chính là thách thức. Các stack HBM4 được đặt cạnh GPU và chỉ có dung lượng giới hạn. Ví dụ, một stack HBM4 12-Hi có dung lượng khoảng 36 GB, trong khi phiên bản 16-Hi đạt khoảng 48 GB.

Điều đó có nghĩa là HBM vừa phải lưu trữ trọng số của mô hình, vừa phải chứa KV Cache. Muốn tăng dung lượng HBM, thông thường phải tăng số lượng GPU, kéo theo chi phí tăng rất nhanh.

HBF có thể thay đổi bài toán này

Như đã đề cập trước đó, SanDisk đang hợp tác với SK hynix để phát triển chuẩn HBF, cung cấp tới 512 GB dung lượng lưu trữ với băng thông từ 0,4 TB/s đến 3 TB/s.

Về nguyên lý, nếu HBM xếp chồng nhiều lớp DRAM thì HBF sẽ xếp chồng nhiều lớp NAND Flash, kết nối với nhau bằng công nghệ Through Silicon Via (TSV) và tích hợp thêm một chip điều khiển (logic die) ở phía dưới.

Điểm yếu lớn nhất của NAND là tốc độ.

SRAM có thời gian đọc khoảng 1 nano giây, DRAM khoảng 100 nano giây, trong khi NAND cần tới khoảng 100 micro giây. Điều này đồng nghĩa với việc tốc độ đọc của NAND chậm hơn DRAM khoảng 1.000 lần.

Để khắc phục hạn chế này, HBF tận dụng khả năng xử lý song song.

Chip điều khiển có thể đồng thời thực hiện hàng nghìn tác vụ đọc trên các ô NAND. Mặc dù mỗi lần đọc vẫn chậm hơn DRAM, nhưng khi hàng nghìn thao tác diễn ra cùng lúc, tổng băng thông đạt từ 0,3 TB/s đến 3 TB/s, cao hơn rất nhiều so với khoảng 6,4 GB/s của chuẩn JEDEC HBM4.

Tuy nhiên, HBF vẫn không thể giải quyết được nhược điểm cố hữu của NAND là tốc độ ghi thấp và độ bền ghi hạn chế.

Vì KV Cache liên tục được cập nhật trong quá trình suy luận AI nên vẫn cần được lưu trên HBM. Ngược lại, trọng số của mô hình vốn ít thay đổi có thể được chuyển sang lưu trên HBF.

HBF không thay thế HBM mà bổ sung cho HBM

Một điểm cần lưu ý là HBF sẽ không có giá rẻ như SSD thông thường vì vẫn phải sử dụng công nghệ đóng gói tiên tiến, vốn có chi phí sản xuất khá cao.

Tuy nhiên, HBF vẫn được kỳ vọng có giá thành thấp hơn đáng kể so với HBM.

Quan trọng hơn, khi HBM không còn phải lưu trữ toàn bộ trọng số của mô hình mà chỉ tập trung cho KV Cache, các hệ thống AI sẽ có thể mở rộng cửa sổ ngữ cảnh dài hơn mà vẫn duy trì hiệu suất và độ chính xác của mô hình. Đây có thể sẽ là một trong những thay đổi quan trọng nhất đối với kiến trúc bộ nhớ dành cho AI trong những năm tới.