Nếu bạn đã từng xem một LLM tạo ra một phản hồi dài và tự hỏi tại sao nó có vẻ chậm dần sau đợt phản hồi nhanh ban đầu, bạn đang quan sát thấy một nút thắt cổ chai phần cứng trong thời gian thực. Hầu hết các nhà phát triển đều đổ lỗi cho mã Python, framework, hoặc kích thước khổng lồ của mô hình. Họ phân tích hiệu năng các hàm, thay đổi bộ tối ưu hóa, và cắt giảm từng mili giây trong quá trình tiền xử lý. Không điều nào trong số đó giải quyết được vấn đề thực sự. Giới hạn tốc độ không nằm ở phần mềm của bạn. Nó nằm ở silicon.
Mọi tác vụ suy luận (inference) của mô hình ngôn ngữ lớn đều phụ thuộc vào hai đặc tính vật lý của GPU trong máy chủ của bạn: tốc độ tính toán số liệu, và tốc độ di chuyển các số liệu đó vào vị trí để tính toán.
Tính toán thì rẻ. Di chuyển dữ liệu thì không.
Tiếp thị GPU rất thích nói về năng lực tính toán (compute). Hàng nghìn tỷ phép tính dấu phẩy động mỗi giây. Những con số đó thật kinh ngạc. Nhưng tính toán chỉ là một nửa câu chuyện. Nửa còn lại là băng thông bộ nhớ (memory bandwidth), tốc độ mà dữ liệu di chuyển từ bộ nhớ băng thông cao vào các lõi tính toán, nơi các phép toán thực tế diễn ra.
Một LLM không thể chạy nhanh hơn liên kết yếu hơn trong hai liên kết này. Hãy tưởng tượng một nhà bếp thương mại với hai mươi đầu bếp bậc thầy. Lò nướng đang nóng, dao đang sắc, và mọi đầu bếp đều sẵn sàng. Nhưng thực phẩm lại được giao bằng xe đạp, từng giỏ một. Nhà bếp bị đình trệ. Thêm đầu bếp sẽ không giải quyết được vấn đề. Mua lò nướng nhanh hơn cũng không giải quyết được. Nút thắt cổ chai nằm ở con đường.
Trong các GPU trung tâm dữ liệu hiện đại, các đơn vị tính toán mạnh mẽ đến mức chúng thường hoàn thành các phép tính của mình rồi ngồi không, lãng phí các chu kỳ trong khi chờ đợi các trọng số (weights) và kích hoạt (activations) truyền qua bộ nhớ. Sự mất cân bằng này không phải là lỗi trong mã của bạn. Đó là thực tế vật lý về cách các con chip được chế tạo. Băng thông bộ nhớ đã không theo kịp năng lực tính toán thô, và các LLM đặc biệt gây khó khăn cho sự mất cân bằng này vì các lượt truyền xuôi (forward passes) của chúng yêu cầu chạm vào mọi tham số cho mỗi token đầu ra.
Tại sao Prompt có vẻ nhanh và Generation có vẻ chậm
Suy luận LLM chia thành hai giai đoạn riêng biệt, và chúng gây áp lực lên phần cứng theo những cách hoàn toàn khác nhau.
Prefill xảy ra khi prompt của bạn vừa chạm tới mô hình. Tất cả các token đều đến cùng một lúc. GPU có thể xử lý chúng song song bằng cách sử dụng các phép nhân ma trận-ma trận lớn. Hàng nghìn đơn vị tính toán hoạt động cùng lúc, và khối lượng công việc luôn dày đặc. Giai đoạn này bị giới hạn bởi tính toán (compute-bound). Sự bùng nổ tốc độ đột ngột mà bạn thấy lúc bắt đầu? Đó chính là lúc GPU đang làm chính xác những gì nó được chế tạo để làm.
Decode là lúc mọi thứ trở nên đau đớn. Khi mô hình tạo ra token tiếp theo, nó thực hiện từng token một. Giai đoạn này dựa trên các phép toán ma trận-vector, vốn chỉ sử dụng một phần rất nhỏ công suất song song của GPU. Tệ hơn nữa, mỗi token mới buộc GPU phải tải lại toàn bộ trọng số mô hình từ bộ nhớ. Các đơn vị tính toán đang muốn làm việc. Thay vào đó, chúng phải chờ đợi. Decode bị giới hạn bởi bộ nhớ (memory-bound). GPU thực chất đang đóng vai trò như một người điều phối giao thông đắt đỏ, vận chuyển các tham số qua lại trên bus bộ nhớ trong khi các công cụ toán học đang "nghỉ ngơi". Đây là lý do tại sao một phản hồi dài một trăm từ có thể mất mười giây mặc dù việc phân tích prompt ban đầu có cảm giác tức thì.
KV cache làm cho điều này trở nên thú vị hơn. Trong quá trình decode, mô hình lưu trữ các tensor key và value cho mọi token trước đó để không phải tính toán lại attention từ đầu. Cache đó tăng dần theo độ dài chuỗi. Nó cũng nằm trong bộ nhớ. Vì vậy, giờ đây GPU không chỉ tải lại trọng số; nó còn đang đọc và ghi một cache đang không ngừng mở rộng trong mỗi lượt truyền xuôi. Các lõi tính toán hầu như không tốn chút sức lực nào trong khi bus bộ nhớ phải làm việc vất vả cho cả hai.
Chống lại "Bức tường bộ nhớ" (Memory Wall)
Các kỹ sư đã phát triển một kho vũ khí nhỏ gồm các kỹ thuật để giảm lượng dữ liệu phải di chuyển, hoặc ít nhất là để chia sẻ chi phí di chuyển chúng.
Batching là cách trực tiếp nhất. Nếu yêu cầu của một người dùng buộc phải tải toàn bộ trọng số từ bộ nhớ, thì việc xử lý tám hoặc mười sáu yêu cầu cùng một lúc cho phép GPU phân bổ tải trọng đó cho tất cả các yêu cầu. Các trọng số được đọc một lần và tái sử dụng cho mọi chuỗi trong batch. Trong môi trường thực tế (production), các hệ thống lập lịch tinh vi sẽ nhóm các yêu cầu một cách linh hoạt, đôi khi được gọi là continuous batching hoặc in-flight batching, để GPU hiếm khi phải tạm dừng. Đó là sự khác biệt giữa một chiếc xe buýt và mười sáu chiếc xe hơi riêng biệt trên cùng một lộ trình.
Quantization giải quyết trực tiếp vấn đề băng thông. Các trọng số mô hình thường được lưu trữ ở định dạng số thực dấu phẩy động 16-bit. Bằng cách nén chúng xuống thành số nguyên 8-bit hoặc thậm chí 4-bit, bạn thực sự cắt giảm lượng dữ liệu truyền qua bus đi một nửa hoặc nhiều hơn. Mô hình vẫn cần đủ độ chính xác để tạo ra đầu ra mạch lạc, nhưng các phương pháp quantization sau huấn luyện (post-training quantization) hiện đại có thể thu nhỏ đáng kể dung lượng bộ nhớ của mô hình mà không làm giảm chất lượng. Ít dữ liệu đang truyền đi hơn đồng nghĩa với việc ít thời gian phải chờ đợi tại bộ điều khiển bộ nhớ hơn.
FlashAttention cấu trúc lại cơ chế attention để giữ các kết quả trung gian bên trong bộ nhớ on-chip tốc độ cao của GPU. Cơ chế attention tiêu chuẩn phải ghi các ma trận attention lớn ra bộ nhớ ngoài chậm chạp rồi sau đó đọc chúng lại. FlashAttention chia nhỏ quá trình tính toán thành các khối (tiles) nhỏ hơn vừa với SRAM, thực hiện các bước softmax và scaling ngay trên chip, và chỉ ghi các kết quả đầu ra cuối cùng trở lại bộ nhớ băng thông cao. Nó đánh đổi một chút tính toán bổ sung để giảm thiểu đáng kể số lần truy cập qua lại tới bộ nhớ chính, và đây gần như luôn là một sự đánh đổi có lợi.
PagedAttention giải quyết một loại lãng phí bộ nhớ khác. Trong quá trình decode, KV cache tăng trưởng một cách không thể dự đoán trước. Các hệ thống truyền thống phân bổ các khối bộ nhớ liên tục và cố định cho mỗi chuỗi, để lại những khoảng trống lớn khi một số chuỗi kết thúc sớm trong khi những chuỗi khác lại mở rộng. PagedAttention mượn khái niệm bộ nhớ ảo từ các hệ điều hành. Nó lưu trữ các mục KV cache trong các khối có kích thước cố định, có thể được phân bổ không liên tục và được ánh xạ thông qua một bảng gián tiếp (indirection table). Điều này ngăn chặn việc bộ nhớ bị bỏ trống trong các bộ đệm đã được dự phòng nhưng chỉ đầy một nửa, đồng thời cho phép kích thước batch lớn hơn, từ đó cải thiện thông lượng tổng thể bằng cách giữ cho bus bộ nhớ luôn bận rộn với các tác vụ hữu ích thay vì lãng phí vào chi phí phân mảnh.
Thay đổi cách đặt câu hỏi
Khi độ trễ tăng vọt, có quá nhiều đội ngũ đặt câu hỏi liệu họ có nên chuyển sang một mô hình nhỏ hơn hoặc viết lại máy chủ suy luận (inference server) hay không. Những câu hỏi đó có quan trọng, nhưng chúng chỉ là thứ yếu. Câu hỏi đầu tiên nên là về chính phần cứng. Liệu GPU của bạn thực sự đang bận tính toán, hay nó đang bị "đói" dữ liệu?
Hãy nhìn vào các chỉ số sử dụng của bạn. Hãy phân tích sự bão hòa băng thông bộ nhớ cùng với mức độ chiếm dụng tính toán của GPU. Nếu bạn thấy sự tranh chấp bộ nhớ cao và cường độ tính toán (arithmetic intensity) thấp trong quá trình decode, bạn không gặp vấn đề về kiến trúc mô hình. Bạn đang gặp vấn đề về vật lý. Giải pháp sẽ không đến từ việc viết code Python sạch sẽ hơn. Nó sẽ đến từ việc batching quyết liệt hơn, quantization các trọng số để lách qua đường ống nhanh hơn, cấu trúc lại attention để giữ nó trên chip, và quản lý KV cache để có thể chứa các batch lớn hơn mà không bị hết bộ nhớ.
Một khi bạn nhìn nhận việc suy luận qua lăng kính này, việc tối ưu hóa sẽ trở nên mang tính cơ học. Bạn sẽ ngừng chạy theo những lầm tưởng rằng trí thông minh của mô hình làm chậm mọi thứ, và bắt đầu đưa ra các quyết định kỹ thuật dựa trên những gì phần cứng thực sự có thể cung cấp. Đó chính là sự chuyển đổi phân biệt giữa các hệ thống sản xuất có khả năng mở rộng và những hệ thống chỉ dừng lại ở mức hoạt động được.