SGLang tối ưu hóa CUDA Graph: Đưa Breakable CUDA Graph làm chuẩn cho Prefill

TL;DR (Tóm tắt)

CUDA Graphs hứa hẹn loại bỏ chi phí khởi chạy kernel (kernel-launch overhead), nhưng để đạt được lợi ích đó trong một công cụ suy luận (inference engine) thực tế, cần phải đồ thị hóa (graphing) khối lượng công việc nhiều nhất có thể mà không làm ảnh hưởng đến tính tương thích, thời gian khởi động hoặc bộ nhớ.

Trong SGLang, chúng tôi đã tái cấu trúc hỗ trợ CUDA Graph xung quanh một giao diện runner/backend chung, giúp các chiến lược capture khác nhau có thể tái sử dụng trên các luồng thực thi. Đối với luồng prefill phức tạp hơn, cộng đồng SGLang đã giới thiệu Breakable CUDA Graph và tiên phong hỗ trợ CUDA Graph đầy đủ với các backend attention FA4 và FlashInfer. Cả hai kỹ thuật này đều được phát triển lần đầu trong SGLang như các kỹ thuật phục vụ mã nguồn mở. Chúng tôi cũng đi sâu hơn vào quản lý bộ nhớ CUDA Graph, bao gồm việc tái sử dụng bộ nhớ trên các hình dạng (shapes) và phân đoạn đồ thị (graph segments), vốn đang trở thành một phần ngày càng quan trọng trong quản lý bộ nhớ tổng thể của SGLang.

Đối với prefill, Breakable CUDA Graph hiện là mặc định của SGLang. Nó đạt được khả năng thực thi phân đoạn tương tự như backend piecewise dựa trên torch.compile với lượng mã nguồn chỉ bằng khoảng một phần tư (521 so với 1.771 dòng), xây dựng các đồ thị prefill nhanh hơn 3,8–5,2 lần vì không cần biên dịch, và có độ bao phủ rộng hơn cho các chức năng phức tạp một cách tự nhiên. Full CUDA Graph cho prefill còn tiến xa hơn, sử dụng request padding để capture toàn bộ quá trình forward ngay cả đối với các khối lượng công việc prefill động. Khi đo lường riêng trên prefill, BCG nhanh hơn 1,70 lần so với thực thi eager (eager execution) và full capture đạt mức 1,93 lần.

Bối cảnh

Một bước suy luận không phải là một kernel đơn lẻ mà là một chuỗi nhiều thao tác GPU. Trong các engine phục vụ LLM hiện đại, việc liên tục khởi chạy các thao tác này từ CPU có thể gây ra chi phí đáng kể, đặc biệt là đối với các khối lượng công việc nhạy cảm với độ trễ. CUDA Graph giảm thiểu chi phí này bằng cách ghi lại công việc GPU một lần và phát lại nó với chi phí khởi chạy thấp hơn nhiều.

Tuy nhiên, việc áp dụng CUDA Graphs hiệu quả trong một engine suy luận hiện đại không hề đơn giản. Thiết kế đồ thị phải phù hợp với các giai đoạn thực thi khác nhau, duy trì tính tương thích với các kernel phức tạp và hành vi phụ thuộc vào runtime, đồng thời kiểm soát thời gian capture và chi phí bộ nhớ do chính các đồ thị tạo ra. Khi các stack suy luận trở nên phức tạp hơn, việc tích hợp CUDA Graph đúng cách ngày càng trở nên quan trọng.

Bài viết này sẽ đi qua cách hỗ trợ CUDA Graph được xây dựng trong SGLang và những thay đổi mà chúng tôi đã thực hiện:

CUDA Graph trong SGLang: Phân tách Runner/Backend và các kết hợp linh hoạt

Trước khi tái cấu trúc này, hỗ trợ CUDA Graph được phát triển xung quanh các luồng thực thi riêng lẻ. Decode, prefill và speculative decoding mỗi loại đều có runner CUDA Graph riêng, với logic chồng chéo về capture shapes, bộ đệm tĩnh (static buffers), phát lại và cấu hình đồ thị. Khi có thêm nhiều chế độ thực thi và chiến lược capture, sự trùng lặp này khiến việc tái sử dụng cơ sở hạ tầng trở nên khó khăn hơn và làm cho các tham số server liên quan đến CUDA Graph ngày càng mơ hồ.

Bản tái cấu trúc [#23906] tách các trách nhiệm này thành hai lớp. Một runner quản lý trạng thái thực thi cụ thể cần thiết cho việc capture và phát lại: các hình dạng đã capture, bộ đệm đầu vào tĩnh, metadata attention và việc đệm (padding) các batch trực tiếp vào các hình dạng đã capture. Một backend xác định cách thực thi đó được capture, cho dù là một đồ thị đầy đủ, một chuỗi các phân đoạn có thể phá vỡ (breakable segments), hay các mảnh do trình biên dịch tạo ra.

Vì các runner chỉ phụ thuộc vào một giao diện backend chung, mỗi luồng thực thi có thể chọn chiến lược capture của riêng mình một cách độc lập. Prefill và decode có các runner riêng biệt, và speculative decoding bổ sung thêm nhiều runner khác: các bước EAGLE draft, draft-extend và frozen-KV MTP draft đều có runner riêng được xây dựng dựa trên decode runner, trong khi target verify chính là decode runner, capture nhiều hơn một token mỗi request.

Runner chuẩn bị cho mỗi luồng thực thi để capture và phát lại, trong khi backend xác định cách quá trình forward được chuyển thành các đồ thị có thể phát lại: dưới dạng một đồ thị đầy đủ, được phân đoạn trong khi capture, hoặc được truy vết (traced) và chia nhỏ trước khi capture.

Full CUDA Graph

Backend đầy đủ capture một torch.cuda.CUDAGraph cho mỗi hình dạng được chọn, không có các vùng eager và có số lần khởi chạy khi phát lại ít nhất trong ba backend. Điều này hoạt động tự nhiên cho decode: mỗi request đóng góp một token, vì vậy biến hình dạng chính là batch size, có thể được bao phủ bởi một tập hợp các bucket batch-size đã capture. Prefill thay đổi theo nhiều chiều hơn và do đó khó hơn; chúng tôi sẽ thảo luận về nó trong phần riêng.

Breakable CUDA Graph

Breakable CUDA Graph (BCG) capture các vùng an toàn với đồ thị (graph-safe) trong khi cho phép các thao tác được chọn chạy ở chế độ eager giữa các phân đoạn đồ thị. Một thao tác không tương thích có thể được đánh dấu bằng @eager_on_graph; quá trình capture dừng lại trước hàm được đánh dấu và tiếp tục sau đó, tạo ra một chuỗi các phân đoạn CUDA Graph được ngăn cách bởi các vùng eager.

Không giống như capture piecewise dựa trên trình biên dịch, các điểm ngắt này được chèn trực tiếp trong quá trình capture thay vì được phát hiện bằng cách truy vết toàn bộ mô hình trước. Chúng tôi thảo luận về cơ chế và lý do tại sao SGLang chuyển sang thiết kế này trong phần tiếp theo.

TC piecewise CUDA Graph

Backend thứ ba đạt được sự phân đoạn tương tự thông qua trình biên dịch. torch.compile truy vết quá trình forward với fullgraph=True, đồ thị FX kết quả được chia tại các điểm chia đã đăng ký, và mỗi mảnh được biên dịch và capture riêng biệt. Đây là câu trả lời đầu tiên của SGLang cho việc capture CUDA Graph một phần và vẫn được sử dụng cho các nền tảng mà capture có thể phá vỡ chưa được xác thực.

Breakable CUDA Graph: Các điểm ngắt Eager không cần trình biên dịch

CUDA Graph theo truyền thống yêu cầu vùng được capture phải hoàn toàn tương thích với đồ thị. Trong thực tế, các khối lượng công việc suy luận hiện đại chứa các thao tác không thể capture trực tiếp. Prefill attention là một ví dụ phổ biến: một số backend attention phụ thuộc vào metadata runtime và sự chuẩn bị phía host. Do đó, một thao tác không tương thích đơn lẻ có thể ngăn CUDA Graph bao phủ một phần lớn hơn nhiều của quá trình forward.

Chúng tôi đã giới thiệu Breakable CUDA Graph (BCG) để làm cho việc capture linh hoạt hơn. Cơ chế và decorator @eager_on_graph lần đầu tiên xuất hiện như một phần của chế độ gỡ lỗi CUDA Graph trong [#19102], và sau đó được tích hợp vào backend piecewise có thể phá vỡ cho prefill trong [#22218]. Thay vì yêu cầu toàn bộ quá trình forward phải tương thích với đồ thị, BCG cho phép các thao tác được chọn chạy ở chế độ eager trong khi vẫn capture các vùng tương thích với đồ thị xung quanh chúng. Ở mức độ cao, quá trình forward trở thành một chuỗi các phân đoạn CUDA Graph được kết nối bởi các điểm ngắt eager rõ ràng.

Thiết kế và Cơ chế

CUDA Graph hoạt động tốt nhất khi việc phát lại tuân theo một chuỗi các thao tác GPU cố định mà không có sự tham gia của host. Tuy nhiên, các quá trình forward suy luận thực tế chứa các thao tác không phù hợp tự nhiên với mô hình đó: các backend attention có thể lập kế hoạch từ độ dài chuỗi trực tiếp, các thao tác tập thể (collectives) có thể liên quan đến sự phối hợp runtime, và các tính năng phục vụ có thể cập nhật trạng thái một cách động.

Việc từ bỏ CUDA Graph bất cứ khi nào một thao tác như vậy xuất hiện sẽ khiến phần lớn quá trình forward không được capture. Thay vào đó, BCG cho phép các nhà phát triển đánh dấu trực tiếp vùng không tương thích bằng @eager_on_graph. Trong quá trình capture, phân đoạn đồ thị hiện tại sẽ đóng lại khi quá trình thực thi đạt đến hàm được đánh dấu, hàm đó chạy ở chế độ eager, và quá trình capture tiếp tục sau đó trong một phân đoạn mới.

Tại thời điểm phát lại, các phân đoạn đồ thị đã ghi và các hàm eager chạy theo cùng một thứ tự. Tensor đi qua điểm ngắt eager được tạo bởi phân đoạn capture trước đó và được đăng ký làm bộ đệm biên (boundary buffer) cố định, vì vậy địa chỉ thiết bị của nó vẫn không đổi. Phân đoạn capture tiếp theo được capture dựa trên cùng địa chỉ đó. Do đó, trong quá trình phát lại, hàm eager ghi kết quả mới tính toán của nó trở lại bộ đệm biên này thay vì trả về một tensor mới được cấp phát, cho phép phân đoạn tiếp theo đọc giá trị đã cập nhật từ địa chỉ mà nó đã được capture ban đầu. BCG không bao giờ kiểm tra hoặc truy vết các thao tác bên trong vùng eager: chúng chỉ cần thực thi chính xác.

Từ góc độ chức năng, BCG và backend piecewise dựa trên torch-compile trước đó tạo ra cùng một loại cấu trúc có thể phát lại: các phân đoạn CUDA Graph được ngăn cách bởi các vùng eager. Sự khác biệt chính là cách cấu trúc đó được xây dựng. TC piecewise trước tiên yêu cầu trình biên dịch hiểu toàn bộ quá trình forward và sau đó chia đồ thị kết quả. BCG đặt các điểm chia trực tiếp trong khi quá trình capture đang diễn ra.

Lợi ích

Khởi động nhanh hơn. Đối với các đồ thị piecewise dựa trên trình biên dịch, việc biên dịch — chứ không phải capture — chiếm ưu thế trong quá trình thiết lập: torch.compile chiếm 78–86% thời gian chuẩn bị các đồ thị prefill, và nó tăng dần theo độ phức tạp của mô hình, đạt 90 giây trên mô hình 235B MoE và 158 giây trên GLM-5.2. BCG loại bỏ hoàn toàn giai đoạn đó, đạt được khả năng thực thi phân đoạn trong một lần capture duy nhất.

Thời gian xây dựng các đồ thị CUDA Graph prefill, 42 hình dạng đã capture, TP4 trên 4×GB300.

Chi phí biên dịch cũng có thể thấy rõ trong quá trình phát triển hàng ngày. Trong thiết lập CI của chúng tôi vào thời điểm đó, việc biên dịch thường được lặp lại qua các lần chạy thử nghiệm, khiến các bài kiểm tra CUDA Graph chậm hơn đáng kể. Việc caching tốt hơn có thể giảm thiểu điều này, nhưng việc loại bỏ trình biên dịch khỏi luồng capture cũng loại bỏ nguồn phức tạp bổ sung này khỏi vòng lặp phát triển.

Tính tương thích rộng hơn. SGLang dựa nhiều vào các kernel tùy chỉnh CUDA, Triton và JIT-compiled không phải là các toán tử PyTorch gốc. Để làm cho các kernel này hiển thị với torch.compile, chúng tôi thường phải bao bọc chúng thông qua torch.library và cung cấp các triển khai giả để truy vết. Điều này đã đưa vào các scaffolding dành riêng cho trình biên dịch trong toàn bộ stack kernel.

Quan trọng hơn, trình biên dịch cũng hạn chế vị trí có thể đặt các ranh giới đồ thị. Các đầu vào và đầu ra đi qua ranh giới toán tử đã đăng ký phải có thể biểu diễn được bởi trình biên dịch. Khi ranh giới tự nhiên liên quan đến trạng thái runtime hoặc kiểu trả về chuyên biệt hơn, đôi khi chúng tôi phải tìm một điểm cắt khác hoặc mở rộng vùng eager chỉ để hiển thị một giao diện mà trình biên dịch có thể xử lý. Khi stack phục vụ phát triển, ranh giới trình biên dịch ngày càng ảnh hưởng đến cấu trúc của mã nguồn vốn không liên quan đến việc biên dịch.

BCG loại bỏ hạn chế này tại các điểm ngắt eager: hệ thống đồ thị không cần hiểu cách hàm được đánh dấu được triển khai hoặc truy vết bên trong nó, cho phép các ranh giới đồ thị tuân theo logic phục vụ thay vì các yêu cầu về truy vết và kiểu dữ liệu của trình biên dịch. Vì CUDA Graph phải cùng tồn tại với DP attention, các backend MoE all-to-all, LoRA, PD disaggregation, hierarchical cache, suy luận tất định và các tính năng phát triển nhanh khác, việc làm cho CUDA Graph hoạt động ngày càng giống như một dự án tích hợp torch.compile. Các kernel mới thường có nghĩa là phải đăng ký custom-op và triển khai giả, trong khi các tính năng mới có thể buộc chúng tôi di chuyển ranh giới đồ thị chỉ để thỏa mãn trình biên dịch. Với BCG, các vùng không tương thích có thể vẫn là thực thi eager thông thường, giảm đáng kể chi phí kỹ thuật dành riêng cho trình biên dịch này.

Có thể gỡ lỗi theo thiết kế. Một CUDA Graph đã capture sẽ phát lại như một đơn vị đóng (opaque unit): Python thông thường không thực thi bên trong nó, điều này làm cho việc in ấn, xác nhận (assertions) và kiểm tra từng bước trở nên khó khăn. BCG tự nhiên để lại các vùng eager nơi Python bình thường vẫn chạy trên mỗi lần phát lại.

SGLang mở rộng ý tưởng này với --debug-cuda-graph [#19102], về cơ bản bao bọc toàn bộ quá trình forward trong một điểm ngắt eager. Mô hình sau đó thực thi ở chế độ eager trong khi vẫn đi qua runner CUDA Graph, bộ đệm tĩnh, luồng phát lại và chuẩn bị metadata. Điều này cung cấp một ranh giới gỡ lỗi hữu ích: nếu vấn đề vẫn còn, nó có khả năng nằm trong mô hình hoặc luồng runner; nếu nó biến mất, bản thân việc capture trở thành nghi phạm chính.

BCG trong Diffusion

BCG cũng đã được áp dụng bởi stack diffusion của SGLang [#27436]. Diffusion liên tục thực thi cùng một quá trình DiT forward trong quá trình khử nhiễu (denoising), làm cho CUDA Graph đặc biệt hữu ích khi các quá trình forward đó chứa nhiều kernel nhỏ, bị giới hạn bởi tốc độ khởi chạy.