Ada 上的 Marlin INT4 — 運作原理

TinyLoop 預設的 INT4 GEMM kernel 是一個簡單的 per-output-element 純量實作（cuda/int2_gemm.cu::int4_gemm）。它正確、小、可攜，但不管 batch size 多少，上限約 6.8 TFLOPS — 因為每個輸出 cell 由單一 thread 計算，權重讀取是頻寬受限，輸出間也沒有重用。

在 Ampere / Ada GPU（SM 8.0 / 8.6 / 8.9）上有個快很多的替代方案：上游 Marlin（Frantar 2024，Apache 2.0），一個張量核心 INT4 × FP16 混合精度 GEMM —— 權重在 global memory 保持 INT4、per-tile dequant 進 shared memory，發 mma.sync 給張量核心。在 RTX 4090 上 尖峰 142 TFLOPS — 大約是 卡本身 FP16 張量核心峰值（≈165 TFLOPS）的 86 %，在 production batch size 下比 naive kernel 快 21 ×。

本頁說明 TinyLoop 如何不改 Marlin 去跑非對稱 GPTQ 權重、runtime 如何自動 dispatch 正確路徑，以及端到端量測長什麼樣。

:::note 硬體範圍上游 Marlin 只支援 Ampere / Ada。在 Hopper（SM 9.0）上 dispatcher 會短路，fallback 回原生 INT4 路徑或 FP16 body cache。詳見上游 README；TinyLoop wrapper 在 dispatch 前呼叫 kernels::marlin_available()，runtime 旋鈕在 kernel 不可用時會自動降級。 :::

核心不匹配：對稱 vs 非對稱

上游 Marlin 的 kernel 假設 對稱 offset-binary 量化：

W_marlin[k, n] = (q_int[k, n] − 8) · scale[g, n]

其中 q_int ∈ [0, 15] 以 4-bit cell 儲存，scale 為 per-group FP16。TinyLoop 的 GPTQ converter（tools/gptq_convert.py::gptq_quantize）產出 非對稱 min-max 權重：

W_tinyloop[k, n] = q_int[k, n] · scale[g, n] + zero[g, n]

多了 per-group-per-column 的 zero offset。把 TinyLoop 權重直接塞給 Marlin 會產出垃圾 —— kernel 隱式減 8 · scale，而 TinyLoop 本來要加 zero。每個 cell 的殘差是：

W_tinyloop − W_marlin = zero + 8 · scale   =:  δ[g, n]

而 matmul 輸出的差值會「塌縮」—— 因為 δ 只依賴 (g, n) 不依賴 k —— 成為一個小小的 per-group 修正：

C_desired[m, n] = C_marlin[m, n] + Σ_g δ[g, n] · (Σ_{k in group g} A[m, k])
               = C_marlin[m, n] + (A_group_sum @ δ)[m, n]

修正是疊在 Marlin 輸出上的一個 [M, groups] × [groups, N] matmul。在我們的 production shape（M ≤ 256, groups ≤ 16, N ≤ 8192）下，約為 Marlin 本身運算量的 ~1 %，主要被 kernel-launch latency 主宰而不是算力。

這正是讓 TinyLoop 能 逐 bit 共用上游 Marlin vendored kernel 檔的原因 —— kernel 內部不需要 zero-point 支援、不需要維護 fork。

端到端 pipeline

                  [載入時，每個 model 一次]
.tinyloop 檔      →  QuantWeight { data, scale, zero, … }          // TinyLoop 原生
                  ↓                                               // (q, scale, zero)
                  ↓  tl.pack_model_for_marlin(model)               // tools/marlin_repack.py
                  ↓                                                //
                  ↓    marlin_pack_asymmetric(q_int, scale, zero)  //
                  ↓      ├─ 重建 w_sym = (q_int − 8) · scale
                  ↓      ├─ 呼叫 marlin_pack_symmetric
                  ↓      │    ├─ torch.round(w_sym / s) → INT4 codes
                  ↓      │    ├─ tile-reshape + permute + _PERM gather
                  ↓      │    └─ 打包 8 × int4 → uint32 cells
                  ↓      └─ δ = zero + 8 · scale    // per group, per column
                  ↓
QuantWeight   ≣  { …,   marlin_B,   marlin_scales,   marlin_delta,   marlin_workspace }

                  [推論時，每次 quant_gemm 呼叫]
quant_gemm(w, X, Y, BT, N, K):
    if marlin_dispatch_enabled()
       && w.has_marlin_pack()
       && marlin_available()
       && BT ≥ 16 && K % 128 == 0 && N % 128 == 0:
        gs = (w.group_size == K) ? -1 : w.group_size      // per-row → per-col
        rc = marlin_mm_asymmetric(                         // include/kernels.h
            X, w.marlin_B, Y, w.marlin_scales, w.marlin_delta,
            w.marlin_workspace, BT, N, K, gs
        )                                                  // 兩個 kernel：
                                                            //   1) marlin_int4_tc_gemm
                                                            //   2) marlin_asymm_correction
        if rc == 0:  return
    // fall through 到原生 INT4 / FP16 body cache 路徑

三個活動元件：

marlin_pack_asymmetric 在 tools/marlin_repack.py — Python 端，載入時透過 torch 在 GPU 上跑。重用 symmetric packer，不重寫 _PERM / _SCALE_PERM permutation。回傳 (B_packed, S_packed, δ) triple。
marlin_asymm_correction 在 cuda/marlin_asymm_correction.cu — 兩個小 CUDA kernel，配一個靜態快取的 scratch buffer：
- a_group_sum_kernel 把 A[M, K] 按 groupsize 大小分塊，reduce 成 a_sum[M, groups]。每個 (m, g) cell 一個 block，128-thread tree reduction。
- correction_accum_kernel 計算 C[m, n] += Σ_g a_sum[m, g] · δ[g, n]，每個 (m, n) 輸出 cell 一個 thread。
Scratch buffer 在呼叫間單調成長（倍增），所以 steady-state 沒有 cudaMalloc 開銷。
marlin_mm_asymmetric 在同一檔 — 把 Marlin 對稱 GEMM + 修正包在一次呼叫裡。先清零 C，回傳 Marlin 的 rc 讓 dispatcher 在 CALL_IF gap 時可以 fallback。

Dispatcher 如何選路徑

src/inference.cpp::quant_gemm 是每個 block 的 attn_qkv / attn_out / mlp_down GEMM 的單一入口（quant_swiglu 裡的 mlp_gate / mlp_up 也是 — 但有個下面會談的 caveat）。分支順序是：

FP16 body cache（w.fp16_cache != nullptr）— 記憶體換速度；在 HBM 頻寬充裕的 Hopper 上勝出。
Marlin + 修正 — 六個前提全部成立時：
- w.bits == 4 && !exact_int4
- marlin_dispatch_enabled() — TINYLOOP_USE_MARLIN=1 或 tl.set_marlin_dispatch(True)
- w.has_marlin_pack() — 載入時有跑 pack
- kernels::marlin_available() — runtime 是 Ampere/Ada
- BT ≥ 16 — 非 per-col shape 的 M 軸 CALL_IF 下限
- K % 128 == 0 && N % 128 == 0 — Marlin tile 大小倍數
Cross-iter stochastic rounding 路徑 — 若 .tinyloop 檔帶有 frac4 或 sign_bit 通道，且 TINYLOOP_CROSS_ITER_ROUND=1 開啟。
原生 INT4 — int4_gemm，永遠可用的 fallback。

步驟 2 中的 (w.group_size == K) ? -1 : w.group_size 轉換值得特別提出。TinyLoop 把 per-row / per-column 量化編碼成 group_size == K（也就是一個大小為 K 的 group）。上游 Marlin 的 CALL_IF 表在 group_blocks > 8 時不實例化 group_blocks = K / thread_k 這個 case —— 它用 groupsize == -1 這個 sentinel 來表示 per-column。這一行翻譯是 pack 與端到端跑通之間最後一個 bug —— 第一版每次呼叫都靜默 fallback 回原生，因為 Marlin 回 rc = 2。

為什麼兩個 kernel 而不是一個融合 kernel

初版融合設計（每個輸出 tile 一個 kernel、讀 A 到 shared + reduce + 累加修正 + 寫 C）在 production shape 下 比較慢。原因：每個輸出 N-tile 都重複對 A[m, :] 做 reduce，N / TILE_N 個 n-tile 就浪費了 ~N / TILE_N × 的 reduction 工作，而兩 pass 設計中 a_sum 只算一次。

目前設計：

每次 Marlin 呼叫兩次 kernel launch（修正用），SM_89 上每個 ≈ 5 µs。
靜態快取 scratch，cudaMalloc 成本攤銷到零。
運算本身低於 1 µs — launch latency 主宰。

可以透過 CUDA Graphs replay 再減半（把 Marlin + 修正 pair 烘進一次圖啟動）。延後實作 —— 它在 Marlin 的 11–66 µs 之上增加的 31–71 % 相對開銷，仍遠低於我們否則卡死在的 6.8 TFLOPS naive 天花板。

RTX 4090 上的端到端結果

tests/test_marlin_e2e.py（一個 CTest target）載入一個真實 .tinyloop INT4 model，在 1 K-token 序列上 A/B Model.score()（dispatch flag 開關），並驗證兩條路徑一致。

在 model_2b_int4_nosvd.tinyloop（per-row INT4，D=2048，FFN=8192，pre=2，loops=8）量測：

路徑	Model.score() PPL（1 K tokens）	備註
原生 INT4（`int4_gemm`）	8.9755	baseline
Marlin + 非對稱修正	8.9157	Δ 0.0597 = 0.67 % PPL drift

15 個權重被 pack（2 pre × 5 linear + 1 loop × 5 linear）
Marlin VRAM overhead：~101 MB
所有 15 個權重每次 forward 都命中 Marlin 分支，零 fallback
首次 dispatch 的 stderr log：[marlin] first dispatch ok: BT=1024 N=6144 K=2048 gs=-1

0.67 % drift 低於整層 × 8 loop iter forward pass 的 FP16 累加雜訊下限。兩份 run 都是同一個量化 model 的有效樣本；沒有誰「比較正確」—— 就像 TINYLOOP_EXPERIMENTAL_FP16_BODY=1 相對 INT4 reference 給出不同但同樣有效的樣本。

Parity 驗證鏈

四個 CTest target 依整合度遞增涵蓋整條 pipeline。全部在 SM 8.9 上 PASS：

Target	層	Shape 數	檢查
`tinyloop.test_marlin_parity`	對稱 packer	3	Marlin vs FP16 ref，僅對稱
`tinyloop.test_marlin_asymmetric_parity`	Python 修正 reference	4	`apply_marlin_correction` 對上 `A @ w_fq`
`tinyloop.test_marlin_asymm_correction_kernel`	CUDA 修正 kernel	5	CUDA 對上 Python ref 在 1–2 ULP 內
`tinyloop.test_marlin_mm_asymmetric`	C++ wrapper 組合	7	`marlin_mm_asymmetric` 在 M=64 N=2048 K=8192（64 groups）以內的 shape
`tinyloop.test_marlin_e2e`	完整 runtime	—	dispatch flag on vs off 的 `Model.score()` parity

五個 test 在前提不成立時都乾淨 skip（exit 77）—— 無 model 檔、Marlin 沒 link、缺 pybind11 等 —— 所以非 GPU 或非 Ada 的 CI runner 維持綠燈。

Gap 與 follow-up

quant_swiglu 尚未 Marlin dispatch。融合 SwiGLU dispatch（一個 kernel 處理 mlp_gate + mlp_up）走自己的路徑，目前還在用 naive INT4 kernel。獨立 follow-up —— 需要 Marlin pair dispatch 或未融合 fallback mode。
Head projection（最後 vocab-size GEMM）未被 pack_model_for_marlin 處理 —— vocab 維度（V = 50257）不能整除 128，所以前提檢查失敗，總之永遠會 fallback。
Hopper（SM 9.0）。Dispatcher gate 已到位但沒有 backend。最乾淨的做法是 vendor gptq-marlin（vLLM fork）—— 它同時有 Hopper WGMMA 路徑跟原生非對稱 scale 支援 —— 但那是遠比 single-file 上游 Marlin 大很多的 vendoring 工作（~數千行 template 重碼）。不在今日 roadmap 上。
修正 kernel pair 的 CUDA Graphs，用來消除小 BT decode 上的 ~10 µs launch-latency 下限。選擇性；目前 overhead 已可接受。

參考

上游 Marlin：github.com/IST-DASLab/marlin — kernel 原始碼 vendored 在 cuda/marlin/，附 Apache 2.0 license。
tools/marlin_repack.py — Python 側 pack + 修正 reference，加上 pack_model_for_marlin driver。
cuda/marlin_wrapper.cu — 我方 namespaced C++ 入口（kernels::marlin_int4_tc_gemm）。
cuda/marlin_asymm_correction.cu — 修正 kernel pair 與組合 wrapper。
Benchmark Comparison — 量測吞吐與精度數字。
環境變數 — TINYLOOP_USE_MARLIN=1 旋鈕文件（見「量化」章節）。

核心不匹配：對稱 vs 非對稱​

端到端 pipeline​

Dispatcher 如何選路徑​

為什麼兩個 kernel 而不是一個融合 kernel​

RTX 4090 上的端到端結果​

Parity 驗證鏈​

Gap 與 follow-up​

參考​