Kimi K3 nhìn từ hai đầu — train trên 72 GPU, chạy trên 8 GB RAM

Tuần này có hai project xuất hiện gần như cùng lúc, ở hai đầu đối lập của cùng một loại kiến trúc. Cursor open-source mixture-of-kittens — một megakernel để train mixture-of-experts trên rack GB300 NVL72. Còn kimi-k3-in-c là bản port C99 để chạy Kimi K3, 2.78 nghìn tỷ tham số, trên một CPU với 8 GB RAM. Đặt cạnh nhau, hai project kể cùng một câu chuyện, và câu chuyện đó không nói về phép nhân ma trận.

Repo: cursor/mixture-of-kittens · FareedKhan-dev/kimi-k3-in-c

Một lưu ý trung thực ngay từ đầu: MoK benchmark trên Kimi K2.7 Code, DeepSeek-V4-Pro, GLM-5.2 và Qwen3.5 — không phải trực tiếp trên K3. Hai bên cùng dòng kiến trúc MoE kiểu DeepSeek-V3, không cùng một checkpoint. Tôi dùng K3 làm trục vì nó là thứ có source đọc được ở cả hai phía.

Bài này đi theo một luận điểm: ở cả hai đầu, kẻ thù chưa bao giờ là FLOPs — mà là việc di chuyển bytes. Và cả hai project, vì hai lý do hoàn toàn khác nhau, đều tự chọn bit-exact determinism làm ràng buộc cứng.

1. Hai con số

MoK, trên 512 GPU GB300 NVL72 trong production của Cursor: 1,070.2 tokens/giây/GPU.

kimi-k3-in-c, trên một CPU với hard cap 8 GB: 32.69 giây/token, peak RSS 8.24 GB, từ một checkpoint 1.56 TB nằm trên NVMe.

Hai con số này không cùng đơn vị công việc — train xử lý hàng nghìn token song song trong một batch, còn decode từng token thì tuần tự và bị chặn bởi latency. Đừng chia hai số cho nhau. Nhưng đặt cạnh nhau chúng đánh dấu hai đầu của một quang phổ, và điều thú vị là lý do mỗi bên chậm lại hoá ra giống nhau đến bất ngờ.

2. Kiến trúc K3

Ba thứ khác thường so với một transformer bình thường:

Attention lai. Không phải mọi layer đều giống nhau. 69 layer dùng KDA — Kimi Delta Attention, một dạng linear attention có recurrence với forget gate theo channel. 24 layer còn lại dùng Gated MLA, nén key/value xuống một latent low-rank. Cứ 4 layer thì 1 layer là MLA, và hai layer cuối đều là MLA để layer cuối cùng luôn làm global attention.

Residual không chạy xuyên suốt. Thay vì cộng dồn một residual qua cả 93 layer, mỗi layer attend lên output của các block trước nó. Layer được chia thành block 12 cái; đến biên block thì residual đang chạy được snapshot lại rồi xoá.

MoE route trong không gian latent. Đường routed chiếu xuống 3584 chiều trước khi dispatch, chứ không dispatch ở full width 7168.

flowchart TD
    E["embed_tokens<br/>163,840 x 7168 bf16"] --> L0
    L0["Layer 0 · KDA + dense FFN 33792<br/>đây là biên block đầu tiên"] --> B1
    B1["Layer 1-3 · KDA + LatentMoE"] --> M4
    M4["Layer 4 · Gated MLA + LatentMoE"] --> DOT
    DOT["... lặp lại: cứ 4 layer thì 1 layer MLA<br/>biên block mỗi 12 layer ..."] --> L92
    L92["Layer 92 và 93 · cả hai đều MLA<br/>layer cuối luôn global attention"] --> N
    N["RMSNorm + AttnRes mức model"] --> H["lm_head<br/>7168 x 163,840"]
    style E fill:#E4ECF3,stroke:#5C7C9B
    style L0 fill:#F2E9D4,stroke:#C0913E
    style M4 fill:#EAE4F1,stroke:#7E6E9C
    style L92 fill:#EAE4F1,stroke:#7E6E9C
    style H fill:#E7ECDD,stroke:#6E9160

3. Một layer chạy như thế nào

flowchart LR
    H["h vào"] --> A1["AttnRes<br/>gộp các block trước"]
    A1 --> ATT["Attention<br/>KDA hoặc Gated MLA"]
    ATT --> A2["AttnRes lần hai<br/>không có guard rỗng"]
    A2 --> M["LatentMoE<br/>hoặc dense FFN ở layer 0"]
    M --> HO["h ra"]
    A1 -.->|"mỗi 12 layer"| SNAP["snapshot<br/>rồi CLEAR"]
    style H fill:#E4ECF3,stroke:#5C7C9B
    style A1 fill:#EAE4F1,stroke:#7E6E9C
    style ATT fill:#F2E9D4,stroke:#C0913E
    style A2 fill:#EAE4F1,stroke:#7E6E9C
    style M fill:#F4E7D3,stroke:#C07A42
    style HO fill:#E7ECDD,stroke:#6E9160
    style SNAP fill:#DDEBE7,stroke:#4E8E86

Hai lần gộp AttnRes không đối xứng: lần trước attention bỏ qua khi chưa có snapshot nào, lần trước MLP thì không có guard đó. Nghe như một chi tiết vặt, nhưng nó quan sát được. In giá trị activation lớn nhất sau mỗi layer, cái reset ở biên block hiện ra rất rõ:

  L45  KDA MoE   |h| max 47.714829
  L46  KDA MoE   |h| max 62.183392
  L47  MLA MoE   |h| max 76.281532
  L48  KDA MoE   |h| max  2.902113   <- biên block, residual bị xoá
  L49  KDA MoE   |h| max  4.353188

Đây là kiểu bug mà một test theo tolerance sẽ cho qua: bỏ cái clear đi thì model vẫn chạy, vẫn sinh ra chữ đọc được, và vẫn là một model khác.

4. MoE — chỗ tốn tiền nhất, ở cả hai đầu

flowchart TD
    X["x · 7168"] --> R["Router · 896 sigmoid ĐỘC LẬP<br/>không softmax, tổng không bằng 1"]
    X --> D["down · 7168 sang 3584"]
    R --> SEL["chọn top-16 theo score + bias"]
    R --> W["trọng số lấy từ score CHƯA cộng bias"]
    D --> Z["z · latent 3584"]
    SEL --> EX["16 expert MXFP4<br/>3584 sang 3072 sang 3584"]
    Z --> EX
    EX --> AGG["cộng có trọng số"]
    W --> AGG
    AGG --> LN["RMSNorm trên TỔNG<br/>không phải từng expert"]
    LN --> UP["up · 3584 sang 7168"]
    X --> SH["2 shared expert · full width"]
    UP --> OUT["out"]
    SH --> OUT
    style R fill:#F2E9D4,stroke:#C0913E
    style EX fill:#F4E7D3,stroke:#C07A42
    style LN fill:#DDEBE7,stroke:#4E8E86
    style OUT fill:#E7ECDD,stroke:#6E9160

Chi tiết dễ làm sai nhất nằm ở router: bias chỉ dùng để chọn, không dùng để tính trọng số. Dùng score đã cộng bias cho trọng số thì vẫn route đến đúng 16 expert đó, chỉ làm lệch tỉ lệ trộn — nghĩa là sai mà không có triệu chứng nào.

Layer này là chỗ tốn tiền ở cả hai đầu. Cursor nói MoE layer “có thể chiếm hơn một nửa thời gian train end-to-end”. Còn ở phía inference, routed experts là 1.45 TB trong tổng 1.56 TB checkpoint — 93%.

5. Đầu train: giấu communication vào bên trong compute

Với expert parallelism, mỗi GPU chỉ giữ vài expert. Token phải được dispatch sang GPU đang giữ expert của nó, tính xong rồi combine ngược về. Đó là hai lần all-to-all mỗi MoE layer, mỗi bước forward và backward. Cursor nói thẳng: “In our production workloads, communication had become the limiting factor.”

MoK trả lời bằng cách gộp tất cả vào một kernel duy nhất:

flowchart TD
    A1["all-to-all dispatch"] --> A2["grouped GEMM"]
    A2 --> A3["SwiGLU"]
    A3 --> A4["grouped GEMM"]
    A4 --> A5["all-to-all combine"]
    A5 --> W["Cách cũ: mỗi mũi tên là một kernel launch<br/>và một lần đồng bộ CPU-GPU"]
    W --> MK["MoK: gộp cả 5 bước vào MỘT megakernel<br/>219 KB CUDA, cả forward lẫn backward"]
    MK --> SM["comp SMs chạy expert FFN<br/>comms SMs đẩy NVLink<br/>hai nhóm báo hiệu nhau qua local counter"]
    style A1 fill:#EAE4F1,stroke:#7E6E9C
    style A5 fill:#EAE4F1,stroke:#7E6E9C
    style MK fill:#F4E7D3,stroke:#C07A42
    style SM fill:#DDEBE7,stroke:#4E8E86

Vài điểm thiết kế đáng chú ý:

• Overlap ở mức SM, không phải mức stream. Một số SM được chỉ định làm compute, số khác làm communication, báo hiệu nhau qua counter cục bộ. Nhờ TMA, NVLink bão hoà được với chưa tới một phần ba số SM.
• Pull cho dispatch, push cho combine. Không đối xứng, và có lý do đo được: pull-based dispatch cho “up to 29% higher NVLink bandwidth utilization” khi expert bị lệch tải, còn signalling là 18 µs so với 103 µs — nhanh hơn 5.8 lần.
• Không còn đồng bộ CPU-GPU. Ring buffer token quay ở mức minibatch, để stream GPU không phải đứng chờ CPU Grace.
• MXFP8, nhưng shared expert giữ BF16 — “as we found this to affect the training stability”.

Kết quả: 2.37× cho MXFP8 forward so với baseline nhanh nhất, và 1.41× end-to-end trên 512 GPU thật, 760.9 lên 1,070.2 tokens/s/GPU.

Một điều cần nói rõ: MoK là CUDA + ThunderKittens, không phải Triton, và chỉ chạy trên NVL72 (Blackwell SM100/SM103). Đây không phải thứ bạn cài lên một node H100.

6. Đầu inference: 134.6 GB mỗi token

Ở phía chạy, con số quyết định mọi thứ là đây:

trunk  (dense, bf16, đọc lại TOÀN BỘ mỗi token)   108.81 GB
routed experts (92 layer x 16 expert x 17.55 MB)   25.83 GB
                                                  ---------
                                                   134.64 GB / token

104 tỷ tham số active, và mỗi token phải kéo tất cả qua bộ nhớ một lần. Với NVMe 3.2–6.1 GB/s, đó là 22–42 giây chỉ để đọc. Compute chiếm khoảng 10 giây trong số đó. Đây là workload bị chặn bởi I/O, không phải bởi ALU.

flowchart LR
    NVME["NVMe · 1.56 TB<br/>1.45 TB routed experts"] --> RING["trunk ring + pinned prefix<br/>108.81 GB bf16"]
    NVME --> LRU["expert LRU · slot 17.55 MB"]
    RING --> TOK["một token<br/>134.6 GB đi qua"]
    LRU --> TOK
    EMB["embed + lm_head + norm<br/>4.70 GB · luôn resident"] --> TOK
    style NVME fill:#ECE8E0,stroke:#8C8477
    style RING fill:#E4ECF3,stroke:#5C7C9B
    style LRU fill:#F2E9D4,stroke:#C0913E
    style EMB fill:#EAE4F1,stroke:#7E6E9C
    style TOK fill:#E7ECDD,stroke:#6E9160

Quyết định kiến trúc quan trọng nhất: expert không bao giờ được dequantize. k3_matmul_mxfp4 nhân thẳng từ nibble đã pack. Một expert là 33,030,144 tham số trong 17,547,264 byte — đúng 0.53125 byte mỗi weight. Nếu bung ra fp32, expert đó thành 132 MB, và một token chạm 1,472 expert: 194 GB mỗi token chỉ để đổi định dạng, trước khi có một phép nhân nào xảy ra.

Cho GPU thêm RAM thì được gì? Đây là bảng thang bộ nhớ, cùng prompt, cùng binary:

Gấp 28 lần bộ nhớ chỉ đổi được 1.7 lần tốc độ. Và cả 12 mức ngân sách đều sinh ra output byte-identical — cùng dãy token id, không sai một bit.

Chú ý cột expert hit: 0.0% từ 8 GB lên tới 64 GB. Tám lần dung lượng, không mua được gì. Giữ ý này lại.

7. Ba sợi dây nối hai đầu

Thang precision

MoK train ở MXFP8, giữ shared expert ở BF16 vì lý do ổn định. K3 ship routed experts ở MXFP4 với quantisation-aware training. Đây không phải cùng một format — đây là hai nấc của một cái thang: train ở 8 bit, ship ở 4 bit, và 4 bit chỉ chịu được vì model đã được train để chịu.

Technical report của K3 nói experts là MXFP4 với QAT, “while all non-expert components remain in higher precision”. Chính câu đó giải thích vì sao 108.81 GB trunk ở phía inference không được phép nén: đo thực tế trên checkpoint thật cho thấy int4 gây ~17% sai số weight trung bình so với ~1% của int8, và những hàng tệ nhất lên tới 65%. Trunk chưa bao giờ được train để chịu 4 bit, nên nó phải được stream nguyên vẹn — đổi giây lấy độ chính xác, vì giây thì mua lại được bằng RAM, còn sai số làm tròn thì không mua lại được bằng gì.

Determinism

MoK: “The order of floating point operations is fixed, so the same input produces bitwise-identical output regardless of hardware scheduling and instruction issue order.” Lý do: ablation nội bộ và on-policy RL post-training.

kimi-k3-in-c đi tới cùng kết luận từ hướng ngược lại. Mọi tích vô hướng cộng dồn bằng double, chia thành 4 accumulator theo i % 4, reduce đúng thứ tự (a0+a1)+(a2+a3). Và đường AVX2 cố tình không dùng FMA:

v = _mm256_add_pd(v, _mm256_mul_pd(wd, xd));   /* NOT fmadd */

FMA làm tròn một lần thay vì hai, nên nó chính xác hơn — và đó chính là vấn đề: nó sẽ cho kết quả khác đường scalar. Một khả năng của phần cứng không được phép làm đổi output.

Hai động cơ khác nhau hoàn toàn — một bên cần RL ổn định, một bên cần output giống nhau ở mọi mức RAM — nhưng cùng dẫn tới một luật: thứ tự phép cộng dấu phẩy động là một phần của hợp đồng, không phải chi tiết cài đặt.

Nghịch lý Quantile Balancing

K3 được train bằng một kỹ thuật tên là Quantile Balancing, mục đích là làm phẳng mức sử dụng expert để không nhóm nhỏ nào chiếm ưu thế. Đó là điều tốt cho model.

Ở phía train, lệch tải vẫn còn đủ để MoK phải có riêng một tham số schedule_capacity_multiplier để chịu đựng nó — khuyến nghị của họ là để cao lúc đầu train khi lệch tải còn tệ, rồi hạ dần về khoảng 0.5.

Ở phía inference, chính sự phẳng đó giết chết LRU cache. Với 16 trên 896 expert mỗi layer và không có nhóm nào nóng, LRU chẳng có gì để giữ. Replay lại trace 100,096 request thật — ở đây trục là dung lượng cache, không phải tổng ngân sách như bảng ở mục 6 — cho thấy LRU đứng yên ở 36.2% từ 8 GB đến 64 GB, trong khi Belady, chính sách tối ưu biết trước tương lai, leo đều từ 39.4% lên 61.7% trên đúng dải đó. Nghĩa là locality có tồn tại, chỉ là LRU không với tới được.

Một quyết định lúc train, sau này, biến thành một dòng 0.0% trong bảng benchmark của người khác.

8. Kết: bandwidth là tường, ở cả hai đầu

Xếp các mắt xích cạnh nhau, với cùng một khối lượng 134.6 GB:

Compute nằm ở đáy bảng, thấp hơn NVMe khoảng ba nghìn lần. Đó là lý do vì sao MoK — dù đã có NVLink 1.8 TB/s trong tay — vẫn phải bỏ 219 KB CUDA ra để giấu communication vào trong compute. Và cũng là lý do vì sao kimi-k3-in-c không tối ưu kernel, mà tối ưu chỗ bytes nằm.

Hai project, hai đầu, một bài học:

• Đừng tối ưu cái không phải bottleneck. Cả hai bên đều đo trước, và cả hai đều tìm thấy movement chứ không phải arithmetic.
• Determinism là một tính năng, không phải sự xa xỉ. Cả hai đều trả giá cho nó, và cả hai đều nói được vì sao đáng.
• Quyết định lúc train sống rất lâu. QAT MXFP4 cho phép 0.53 byte/weight lúc chạy; Quantile Balancing thì biến cache thành vô dụng. Không cái nào là bug.

Nếu bạn định port thứ này lên GPU, con số cần nhìn không phải TFLOPS — mà là có bao nhiêu trong 134.6 GB kia bạn có thể ngừng di chuyển.