SGLang và Miles hỗ trợ tối ưu hóa mô hình Kimi K3 2.8T tham số ngay khi ra mắt

Chúng tôi rất vui mừng thông báo về việc hỗ trợ Day-0 cho Kimi K3 trong SGLang và Miles. K3 là mô hình mã nguồn mở đầu tiên thuộc phân khúc 3 nghìn tỷ tham số, và kiến trúc lai của nó đi chệch khỏi các quy ước thông thường ở hầu hết mọi nơi mà một ngăn xếp phục vụ (serving stack) thường mặc định. Với sự hợp tác cùng các đội ngũ Moonshot AI và NVIDIA, cả hai công cụ này đều hỗ trợ đầy đủ cho K3 ngay trong ngày ra mắt: SGLang cho suy luận (inference), Miles cho huấn luyện RL. Bài viết này sẽ đề cập đến những công việc cần thiết để đạt được điều đó.

Các điểm nổi bật

Các lệnh khởi chạy và hướng dẫn cấu hình theo khối lượng công việc (workload) đều có sẵn trong Kimi K3 cookbook.

Triển khai Kimi K3

Kimi K3 là mô hình mã nguồn mở đầu tiên thuộc phân khúc 3 nghìn tỷ tham số: 2,8 nghìn tỷ tham số, cửa sổ ngữ cảnh 1 triệu token và khả năng hiểu hình ảnh gốc. Người tiền nhiệm K2.5 của nó có kiến trúc khá gần gũi với các mô hình mà SGLang đã phục vụ tốt, vì vậy hầu hết các thành phần trong ngăn xếp đều có thể áp dụng trực tiếp. K3 thì không như vậy — nó đi chệch khỏi các quy ước thông thường ở nhiều khía cạnh độc lập cùng một lúc:

Mỗi hàng trong số này đều là một bài toán phục vụ. Ngăn xếp attention lai (hybrid attention stack) là yếu tố giúp cửa sổ ngữ cảnh 1 triệu token trở nên khả thi, nhưng nó đồng nghĩa với việc máy chủ phải lưu giữ hai loại trạng thái cùng lúc: một trạng thái KDA kích thước cố định cho mỗi yêu cầu bên cạnh KV theo từng token của MLA — các phần quản lý bộ nhớ bên dưới chính xác là nói về điều này. Attention Residuals luân chuyển một tập hợp các đầu ra attention qua toàn bộ ngăn xếp, điều này phá vỡ các giả định mà hệ thống layer tiêu chuẩn của SGLang đặt ra và buộc phải tạo các đường dẫn riêng cho K3 ở những nơi như DP attention. LatentMoE định tuyến 16 trong số 896 chuyên gia (experts) bên trong không gian tiềm ẩn (latent space) được chiếu xuống với hàm kích hoạt SiTU thay vì SwiGLU, vì vậy không có kernel MoE hiện có nào có thể áp dụng ngay lập tức. Và đường dẫn thị giác — bao gồm tower, projector, processor và định dạng phương tiện XTML của Kimi — đã được xây dựng từ đầu. Các phần bên dưới đều dựa trên quá trình triển khai này.

Quản lý bộ nhớ KDA lai

Attention KV chỉ có thể được thêm vào (append-only). Khi KV của một token được tính toán, nó sẽ không bao giờ thay đổi, điều này cho phép trình lập lịch (scheduler) chia sẻ một bản sao vật lý cho mọi yêu cầu có chung tiền tố, giữ nó trong một radix tree và thực hiện pipeline qua các lần lặp mà không cần suy nghĩ thêm. Trạng thái của một layer KDA thì ngược lại, đó là một bộ đệm tái phát (recurrent buffer) kích thước cố định bị ghi đè tại chỗ ở mỗi token. Vì vậy, mọi thứ mà trình lập lịch cung cấp cho KV một cách miễn phí, từ prefix caching đến overlap scheduler, speculative decoding cho đến paging, đều phải được xây dựng lại cho một giá trị thay đổi ngay khi nó được đọc. K3 xen kẽ 69 layer KDA với 24 layer MLA, vì vậy đây không phải là một trường hợp ngoại lệ. Đó là phần lớn mô hình. Những gì theo sau là quá trình xây dựng lại đó, và đặc điểm thỏa đáng duy nhất của nó là mọi thành phần đều bị ép buộc bởi thực tế ghi đè tại chỗ. Không có gì trong số đó là các thành phần chắp vá.

Một đường dẫn yêu cầu duy nhất

Cùng một đường dẫn phục vụ cho overlap scheduler, speculative decoding và page_size > 1 cùng lúc, những sự kết hợp vốn trước đây loại trừ lẫn nhau đối với các mô hình trạng thái tái phát. Một yêu cầu có thể truy cập vào tiền tố đã lưu trong bộ nhớ đệm, khôi phục trạng thái tái phát, chạy một bản nháp đa token, xác minh nó và cam kết (commit), tất cả trong khi trình lập lịch chuẩn bị lô tiếp theo trước một bước.

Ba bước chuyển trạng thái trên luồng forward

Trạng thái trực tiếp của một yêu cầu nằm trong một khe làm việc (working slot) duy nhất, được đọc và ghi đè tại chỗ bởi mỗi lần forward. Việc lưu vào bộ nhớ đệm (caching) có nghĩa là sao chép ra khỏi bộ nhớ mà GPU đang tích cực thay đổi, và điều đó tạo ra hai cuộc đua (race). Một quá trình khôi phục có thể va chạm với lần forward đang ghi vào khe đó. Và ảnh chụp nhanh (snapshot) làm mới bộ nhớ đệm có thể va chạm với một lệnh donate vẫn đang đọc ảnh chụp trước đó. Cả hai đều kết thúc mà không cần đồng bộ hóa toàn thiết bị và không có khóa trên đường dẫn nóng (hot path).

Thứ nhất, mỗi bản sao trạng thái là một kernel trên luồng forward tuần tự. Quá trình copy-on-write khôi phục một checkpoint đã lưu vào khe làm việc, và snapshot ghi lại trạng thái tại ranh giới track, được xếp hàng giữa các lần forward tạo ra và tiêu thụ trạng thái đó, vì vậy thứ tự trên cùng một luồng cung cấp cơ chế happens-before miễn phí. Một snapshot được thực hiện một lần cho mỗi khối prefill và một lần mỗi khoảng track trong quá trình decode.

Thứ hai, bước chuyển duy nhất không để lại yêu cầu nào truyền tải byte. Các snapshot nằm trong một cặp ping-pong mà chúng tôi gọi là bộ đệm phụ (extra buffer), nơi một khe giữ snapshot mới nhất trong khi khe kia nhận snapshot tiếp theo, và khe thứ hai đó chỉ được cấp phát tại ranh giới cần nó và giải phóng ngay sau đó. Caching sẽ donate chỉ số khe của snapshot mới nhất vào cây, sau mỗi khối prefill, tại thời điểm bàn giao từ prefill sang decode, và khi yêu cầu kết thúc. Một khe mới sẽ lấp đầy bộ đệm, vì vậy mục tiêu ghi của snapshot và nguồn đọc của lệnh donate không bao giờ là cùng một khe vật lý.

Ba bước chuyển trạng thái: Copy-on-write, snapshot và donate, cùng vị trí của mỗi bước so với luồng forward tuần tự.

Khe theo yêu cầu, checkpoint dùng chung

Vì các checkpoint tiền tố có thể tái sử dụng nằm trong radix tree dùng chung và có thể bị loại bỏ (evictable), một yêu cầu đang chạy chỉ dành riêng một vài khe tạm thời, ít nhất là bốn. Đó là khe làm việc của nó, một khe bộ đệm phụ cho các snapshot, và hai khe dự phòng lưu giữ, một cho trạng thái tiền tố đã cam kết cần phải duy trì và một cho bản sao mà một nhánh phân kỳ cần. Phần lớn trạng thái được lưu trong bộ nhớ đệm được phân bổ trên toàn bộ cây thay vì tính phí cho mỗi yêu cầu, vì vậy nhóm trạng thái không cần phải tăng lên theo dung lượng lịch sử mà mỗi yêu cầu giữ lại.

Prefix caching cho trạng thái tái phát

Một trạng thái tái phát không thể bị cắt tại một token tùy ý, vì bạn không thể chạy ngược nó về vị trí trước đó, nên nó chỉ được checkpoint tại các ranh giới khối, và ngay cả ở đó cũng rất thưa thớt. Một giới hạn theo đường dẫn và LRU giữ cho chỉ một vài checkpoint tồn tại trên mỗi đường dẫn, và một checkpoint có thể bị loại bỏ độc lập với KV mà nó chú thích, để lại nút đó như một bia mộ. Các điểm phân nhánh được xử lý đặc biệt, một ý tưởng từ Marconi. Một nhánh rẽ là tiền tố duy nhất mà mọi nhánh tương lai chắc chắn sẽ chia sẻ, vì vậy khi một yêu cầu phân kỳ giữa chừng, nó sẽ phát lại từ checkpoint gần nhất phía trên và đặt một checkpoint mới tại điểm phân nhánh đã căn chỉnh theo khối. Nhánh tiếp theo sẽ khôi phục trực tiếp tại đó, không cần phát lại.

Các checkpoint trên radix tree. Lớp phủ checkpoint thưa thớt và điểm phân nhánh.

Không có cái nào trong số này là một tính năng riêng biệt được chắp vá vào các tính năng khác. Chúng là bốn hệ quả của một sự thật: trạng thái tự ghi đè lên chính nó, đó là lý do tại sao thiết kế vẫn nhỏ gọn ngay cả khi nó tích hợp đồng thời overlap scheduler, speculative decoding, paging và prefix caching.

Bộ nhớ thống nhất: một nhóm cho cả hai loại trạng thái

Mọi thứ ở trên quản lý từng loại trạng thái bên trong nhóm riêng của nó, và bản thân các nhóm này là phỏng đoán duy nhất còn lại trong thiết kế. Hai đơn vị cấp phát cách nhau ba bậc độ lớn: một khối trạng thái KDA lớn cho mỗi yêu cầu (khoảng 54 MB với TP=8, bao phủ tất cả 69 layer KDA) và một khối MLA KV nhỏ cho mỗi token (khoảng 27 KB, bao phủ tất cả 24 layer MLA), vì vậy hiện tại chúng nằm trong hai nhóm riêng biệt, được định kích thước khi khởi động. Việc định kích thước đó là một sự đặt cược vào lưu lượng truy cập, và khi sự đặt cược đó sai, máy chủ sẽ hết bộ nhớ ở một nhóm trong khi nhóm kia vẫn còn dư.

Bộ nhớ thống nhất thay thế hai nhóm bằng một: các trạng thái KDA lấp đầy từ một đầu, các khối MLA KV từ đầu kia, và các byte không sử dụng ở giữa tạo thành một vùng trống duy nhất.

Bộ nhớ thống nhất. Trên: hiện tại hai loại trạng thái có các nhóm riêng biệt, được định kích thước khi khởi động, vì vậy một nhóm có thể nhàn rỗi trong khi nhóm kia đầy. Giữa: với bộ nhớ thống nhất, cả hai loại đều cấp phát từ cùng một nhóm, phát triển từ các đầu đối diện với một vùng trống duy nhất ở giữa. Dưới: việc giải phóng một khối trạng thái ở giữa để lại một khoảng trống, và một khối từ cuối được di chuyển vào đó để vùng trống luôn là một khối liền mạch.

Việc giải phóng cũng đơn giản như vậy. Khi một yêu cầu kết thúc, bị hủy bỏ hoặc bị thu hồi dưới áp lực, các khối của nó sẽ được giải phóng. Nếu điều đó để lại một khoảng trống ở giữa, một khối từ cuối được di chuyển vào đó để vùng trống luôn là một khối liền mạch.

Bố cục này mang lại cho chúng tôi kích thước trang linh hoạt và không gây phân mảnh bộ nhớ: một khối trạng thái KDA 54 MB và một khối MLA KV 27 KB lấy từ cùng một byte mà không bị ép buộc vào bất kỳ kích thước trang chung nào, và các bước di chuyển ở trên giữ cho mỗi đầu được đóng gói với chi phí không đáng kể cho việc di chuyển trạng thái, vì vậy không gian trống luôn là một vùng liền mạch có thể sử dụng bởi cả hai loại. Do đó, dung lượng tuân theo khối lượng công việc thay vì một cờ khởi động: nhiều yêu cầu ngắn lấp đầy nhóm bằng các khối trạng thái, một vài ngữ cảnh dài lấp đầy nó bằng KV, và không trường hợp nào cần cấu hình lại bất cứ thứ gì.

Bộ nhớ thống nhất được phát hành dưới dạng tùy chọn thông qua --enable-unified-memory. Một bài viết tiếp theo sẽ đi sâu vào chi tiết triển khai.

Speculative decoding với DSpark

K3 đi kèm với DSpark block speculative decoding, được điều khiển bởi một mô hình nháp (draft model) mà chúng tôi đã huấn luyện cho K3. Hai phần của quá trình tích hợp xứng đáng có một câu chuyện riêng: chỉ chi tiêu ngân sách xác minh ở nơi mang lại hiệu quả, và làm cho trạng thái KDA tái phát của K3 có thể tồn tại được với quá trình suy đoán (speculation).

Chỉ xác minh những gì đáng xác minh

DSpark đề xuất một khối các token nháp cho mỗi bước, và mô hình mục tiêu xác minh toàn bộ khối trong một lần forward. Ở batch size 1, các vị trí xác minh bổ sung về cơ bản là miễn phí: bước này bị giới hạn bởi độ trễ (latency-bound), và thêm một vài token nữa cũng không ảnh hưởng. Khi batch đầy, điều đó không còn đúng nữa. Các token xác minh giờ đây cạnh tranh với mọi yêu cầu khác cho cùng một thời gian bước, và hầu hết chúng đều thua cược: trên khối lượng công việc chat, độ dài chấp nhận là khoảng 2,7, vì vậy năm trong số tám vị trí được xác minh trên một bước điển hình bị từ chối. Xác minh mọi thứ đồng nghĩa với việc trả giá đầy đủ cho các token mà máy chủ sau đó sẽ loại bỏ.

Các mảnh ghép để sửa lỗi này đã có sẵn trong hệ thống. Bản nháp mang theo một confidence head đã được huấn luyện để dự đoán, theo từng vị trí, khả năng mỗi token vượt qua xác minh là bao nhiêu. Ở phía bên kia, một hồ sơ đo lường một lần của máy chủ ghi lại chi phí thực tế của một token xác minh bổ sung ở mỗi mức tải. Một trình lập kế hoạch theo từng bước kết hợp cả hai: mỗi yêu cầu chỉ giữ lại các token xác minh khi giá trị kỳ vọng của chúng bù đắp được chi phí biên, và phần còn lại của cửa sổ sẽ bị cắt tỉa trước khi quá trình forward mục tiêu khởi chạy. Những gì còn lại được xác minh chính xác như trước, vì vậy đầu ra vẫn không bị mất dữ liệu. Sự đánh đổi là một chuỗi chấp nhận ngắn hơn một chút để đổi lấy một bước rẻ hơn.

Cắt tỉa mang lại hiệu quả khi chịu tải. Thông lượng giải mã (decode throughput), xác minh tất cả so với cắt tỉa, trên bảng chat (chấp nhận ~2,7, bên trái) và bảng toán học few-shot (chấp nhận ~5,0, bên phải). Hòa vốn đến bs 8, sau đó khoảng cách mở rộng theo batch size: +68% và +24% tại bs 256, với độ dài chấp nhận giảm từ 2,7 xuống 2,2 và từ 5,0 xuống 4,3 trong khi thông lượng tăng lên.

Đường cong chi phí đo lường hóa ra lại là phần thú vị. Chi phí biên của một token xác minh không mượt mà. Nó là một cầu thang: các kệ phẳng nơi một token khác lọt vào các sóng kernel hiện tại với chi phí gần như bằng không, và các bậc thang dốc nơi nó bắt đầu một sóng mới. Trình lập kế hoạch đọc bề mặt này. Khi các token tiếp theo của một yêu cầu nằm trên một kệ rẻ, nó giữ lại chúng; khi chúng bắt đầu một bậc thang, nó cắt tại đó. Trong dữ liệu, điều này xuất hiện dưới dạng các dao động nhỏ trong đường cong độ dài chấp nhận trong khi thông lượng vẫn mượt mà và đơn điệu: trình lập kế hoạch đang lướt trên các kệ phần cứng thực tế, không phải nhiễu.

Dưới batch size 8, không có nhiều áp lực cần giải tỏa, vì vậy việc cắt tỉa chỉ ở mức hòa vốn đến hơi tiêu cực; một lối thoát sớm cho batch nhỏ trong trình lập kế hoạch là phần tiếp theo đã được xác định.

ReplaySSM: phát lại đầu vào thô cho trạng thái KDA

Speculative decoding xác minh γ+1 token nháp cùng một lúc và có thể chỉ chấp nhận một tiền tố. Đối với MLA, điều đó là miễn phí, vì KV chỉ có thể thêm vào và một bản nháp bị từ chối chỉ giải phóng các khe của nó. Trạng thái của một layer KDA tự ghi đè lên chính nó ở mỗi token, như phần trước đã đề cập, vì vậy cơ sở này mua khả năng đảo ngược bằng vũ lực và chụp ảnh nhanh toàn bộ trạng thái K×V sau mỗi bước nháp. Tại K=V=128, đó là 64 KB cho mỗi yêu cầu, mỗi layer và mỗi head, nhân với γ+1 bước. Trên 69 layer KDA của K3 và một batch đầy đủ, nó vượt quá nhóm trạng thái bền vững mà nó đang cạnh tranh, và vì nó được dành riêng cho mỗi yêu cầu đang chạy, nó giới hạn khả năng đồng thời.

Lưu trữ đầu vào, không phải trạng thái. ReplaySSM loại bỏ các snapshot. Kernel xác minh đọc checkpoint đã cam kết và không bao giờ ghi vào nó, và trong quá trình thực hiện, nó cũng lưu trữ các đầu vào thô của mỗi bước Sᵢ = (vᵢ, kᵢ, gkᵢ, βᵢ), khoảng 1 KB so với 64 KB mà một snapshot tiêu tốn. Khi bộ lấy mẫu (sampler) xác định độ dài chấp nhận, một kernel fold duy nhất bao phủ mọi layer và head sẽ phát lại chỉ tiền tố được chấp nhận từ checkpoint và tiến hành nó tại chỗ. Các bản nháp bị từ chối không bao giờ được phát lại, vì vậy chi phí rollback bằng không. Cửa sổ nháp giảm từ 512 KB xuống 16 KB, khoảng 32 lần.