昨天看了 SGLang RVV 後端怎麼從 Python 接到 C++。今天先來講 linear / GEMM(General Matrix Multiplication)，因為 LLM 裡的 QKV projection、MLP gate/up/down projection、lm_head 都會大量使用矩陣運算。

今天沿著這條路往下看

python/sglang/srt/layers/rvv_utils.py
-> torch.ops.sgl_kernel.convert_weight_packed
-> torch.ops.sgl_kernel.weight_packed_linear
-> sgl-kernel/csrc/cpu/riscv64/gemm.cpp

本篇大綱

• 先看 LLM 裡哪些位置會呼叫 linear
• 再看為什麼 CPU 後端需要 weight packing
• 接著讀 weight_packed_linear 的 Python 分派 和 C++ 註冊
• 然後拆 RVV GEMM迴圈：vsetvl、BF16/FP16 載入、FP32 累加、儲存
• 最後回到 K1 的 Vector-256bit、L1 tiling、prefetch 與 OpenMP，看這些最佳化如何配合

Linear 在 LLM 裡出現在哪裡

Transformer 裡的 linear 很多。以 decoder-only LLM 來看，一層裡通常會碰到

attention:
  q_proj / k_proj / v_proj / o_proj

MLP:
  gate_proj / up_proj / down_proj

最後輸出:
  lm_head

這些 linear 都在把 token 目前的 hidden state 投影到另一個特徵空間。每個投影使用不同的 weight，所以學到的功能也不同。

• Attention 裡的 q_proj、k_proj、v_proj 分別產生 query、key 和 value。Query 和 key 用來計算 token 之間的關聯分數，再依分數對 value 加權求和。o_proj 會把各個 attention head 的結果混合後投影回 hidden size，讓結果可以接回 residual path。Q/K/V projection 是在準備 attention 的輸入，後面還會有 QK^T 內積、softmax 和對 V 的加權求和。
• MLP 裡的 gate_proj 和 up_proj 會把 hidden state 擴展到較大的中間層維度。以 SwiGLU 為例，gate_proj 的輸出先經過 SiLU，再與 up_proj 的輸出逐元素相乘；down_proj 再把結果壓回 hidden size。這一段讓每個 token 獨立地進行非線性特徵轉換。
• 模型最後的 lm_head 把 hidden state 投影到 vocabulary size。輸出的每一個數是對應 token 的 logit，後續再經過 softmax 或取樣決定下一個 token。

這些 linear 的基本形式都可以寫成：

Y = X @ W^T + bias

在推論服務裡，prefill 和 decode 的形狀很不一樣。

Prefill 是讀取 prompt 的階段。模型會一次處理多個 token，並為它們建立後續生成會用到的 KV cache，所以 M 通常大於 1。例如 M=64、K=2048、N=3072 的 QKV projection。

Decode 是逐步產生新 token 的階段。每一步會重用先前的 KV cache，只對最新的 token 進行模型計算。單一請求下 M 常常是 1，這時 linear 更像 GEMV（General Matrix-Vector Multiplication）：

1 x K activation
K x N weight
-> 1 x N 輸出

CPU 上 linear 的瓶頸

在 GPU 上，GEMM 通常有成熟的函式庫和很高的記憶體頻寬。RISC-V CPU 上，尤其是 Banana Pi / SpacemiT K1 這類板子，linear 會受幾個因素影響

• RVV 每輪能處理多少元素
• weight資料配置是否適合向量載入
• activation 和 weight 是否連續
• L1/L2 cache 能不能留住重複使用的資料
• LPDDR 頻寬是否已經成為瓶頸
• 多核心 執行緒切分是否造成頻寬競爭

前面提到 Banana Pi BPI-F3 的 Vector-256bit，如果用 e32m1 估算，一輪最多處理 8 個 FP32。
如果 BF16 載入後 widen 到 FP32，常見寫法會用 BF16 載入、轉成 FP32 累加器，再做 vfmacc。

linear 的效能通常不只看算術指令，weight 每一層都很大，如果資料配置讓 kernel 一直做不連續讀取，RVV FMA 指令再漂亮也吃不到完整效果。

Weight packing 解決什麼問題

一般 PyTorch weight 常見形狀是

weight: [N, K]

如果 kernel 要同時算一個 N 方向分塊，它希望 weight 在記憶體裡排成比較適合向量載入的形狀。目前 K1 開發分支中這條 BF16 / FP16 packed linear 的 BLOCK_N 是 64；block size 是實作參數，不應從 256-bit VLEN 直接推成 32。

SGLang RVV 後端會在載入 weight 之後做 packing

def _rvv_process_weight_after_loading(module, weight_names) -> None:
    ...
    convert_weight_packed = _get_convert_weight_packed_op()
    ...
    for name in weight_names:
        w = getattr(module, name)
        packed = torch.nn.Parameter(convert_weight_packed(w.data), requires_grad=False)
        packed.__dict__.update(w.__dict__)
        setattr(module, name, packed)

    module.use_riscv_rvv_backend = True

這裡補充一下 convert_weight_packed 不是對所有 unquantized weight 都產生同一種 blocked layout

1. 確認 weight 在 CPU 上。
2. 確認機器支援 RVV。
3. 當輸出通道數適合 blocked kernel 時，產生 [NB, K * BLOCK_N] 的 packed weight。
4. 對於非 INT8、輸出通道較小或無法被 block 整除的形狀，目前實作會回傳 FP32 transpose，後續走 FMA fallback。

這個預處理的核心想法是把每次 kernel 都要重排資料移到模型載入後做一次。後面每次前向傳播直接使用載入時產生的 representation：它可能是 blocked packed weight，也可能是 FP32 transpose。

Python 分派 怎麼走到 packed linear

在 unquantized linear 路徑裡，SGLang 會檢查 layer 是否已經走 RVV 後端

if use_intel_amx_backend(layer) or use_riscv_rvv_backend(layer):
    x_shapes = x.shape
    if len(x_shapes) == 3:
        x = x.view(-1, x.shape[-1])
    with torch.profiler.record_function(
        "sglang.rvv.linear.weight_packed_linear"
    ):
        output = torch.ops.sgl_kernel.weight_packed_linear(
            x,
            layer.weight,
            bias,
            True,  # is_vnni / is_packed
        )
    if len(x_shapes) == 3:
        output = output.view(x_shapes[0], x_shapes[1], -1)
    return output

這裡有幾個細節

第一，3D 輸入會先 flatten 成 2D。很多 linear kernel 的 C++ 實作比較喜歡 [M, K] 形狀。

第二，layer.weight 已經在載入時被換成 convert_weight_packed 的結果，可能是 blocked layout 或 FP32 transpose。

第三，is_packed=True 在這個 API 中表示 weight 已經預處理，C++ 不再呼叫 converter，它不保證 representation 一定是 blocked layout。

第四，效能分析器標記會出現 sglang.rvv.linear.weight_packed_linear。後面做歸因時，可以用效能分析器或追蹤檢查熱路徑是否真的走到這裡。

C++ 端怎麼接住

weight_packed_linear 在 sgl-kernel/csrc/cpu/riscv64/gemm.cpp。它會先處理形狀、資料型別、bias、packed資料配置檢查：

auto packed_w = is_packed ? mat2 : convert_weight_packed(mat2);
...
int64_t M = mat1.size(0);
int64_t K = mat1.size(1);
...
auto out = at::empty({M, N}, mat1.options());

如果 is_packed=False，它會當場呼叫 convert_weight_packed(mat2)。這可以讓 API 彈性更高，但推論服務熱路徑通常希望 weight 早就 pack 好。

接著它會區分 FP32-transpose FMA fallback 和 blocked-packed kernel。只有 blocked-packed 分支要求 activation 與 weight 同 dtype：

if (!use_fma_gemm) {
  TORCH_CHECK(
      packed_w.scalar_type() == mat1.scalar_type(),
      "weight_packed_linear: mat1 and weight must have the same dtype ...");
}

目前路徑會透過 AT_DISPATCH_RVV_TYPES分派；這個 macro 也包含 FP32，不只是 BF16 / FP16：

AT_DISPATCH_RVV_TYPES(dispatch_type, "weight_packed_linear_kernel_impl", [&] {
  if (use_fma_gemm) {
    weight_packed_linear_fma_kernel_impl<scalar_t>(
        out.data_ptr<scalar_t>(),
        mat1.data_ptr<scalar_t>(),
        packed_w.data_ptr<float>(),
        ...);
  } else {
    weight_packed_linear_kernel_impl<scalar_t>(
        out.data_ptr<scalar_t>(),
        mat1.data_ptr<scalar_t>(),
        packed_w.data_ptr<scalar_t>(),
        ...);
  }
});

這裡的 scalar_t 可以是 FP32、BF16 或 FP16。對 BF16 / FP16 路徑，kernel 內部會轉成 FP32 累加，最後再存回輸出資料型別。

RVV GEMM迴圈 的形狀

GEMM 內部可以先用簡化版理解

for m block:
  for n block:
    acc = 0
    for k:
      acc += A[m, k] * packed_W[n_block, k]
    store C[m, n]

RVV 向量化維度要以實作為準。目前 gemm.cpp 的 BF16 / FP16 packed linear 是沿 N（輸出通道）向量化：

1. K 迴圈每次讀一個 activation A[m, k]，作為 scalar 廣播。
2. 從 packed weight 連續載入一段 N 方向資料。
3. BF16 / FP16 weight 轉成 FP32 向量。
4. 用 vfmacc_vf 將 scalar activation 與 weight vector 累加到 FP32 向量累加器。
5. 每個 K tile 結束後，把 FP32 partial sum 寫到對齊的 C_f32 暫存區；所有 K tile 完成後，再轉回輸出型別。這條路徑沒有 K-vector horizontal reduction。

linear / GEMM 多了 K-迴圈和累加器。這份實作會讓一個 K tile 內的累加器留在向量暫存器，tile 之間再透過 C_f32 保留 partial sum。其他重點還包括 weight 載入能不能連續，以及 K-迴圈能不能減少多餘的轉換和記憶體流量。

BF16 / FP16 / INT8 / INT4 路徑

這篇主要看 BF16 / FP16 packed linear，但 SGLang RVV 後端還有量化路徑。

W8A8 INT8 會走 int8_scaled_mm_cpu 或 int8_scaled_mm_with_quant。這類 kernel 需要處理 activation / weight 縮放係數，並且把 int8 multiply 累加轉成目標資料型別。

W4A8 動態 INT4 會在載入時呼叫：

convert_weight_w4a8_dynamic_packed(weight_data, scales_data, group_size)

它會產生：

module._rvv_int4_w4a8_dynamic_w_q
module._rvv_int4_w4a8_dynamic_w_s
module.use_riscv_rvv_int4_w4a8_dynamic_linear_backend = True

量化路徑的工程問題會更多：unpack、縮放係數資料配置、group size、act-order、輸入 permutation、lm_head tied embedding。這些都會影響 decode 延遲。先把 BF16 packed linear 看懂，後面再看 INT8 / INT4 比較不會混在一起。

這份 GEMM 怎麼利用 K1 的 CPU 特性

回到 gemm.cpp，這份 linear kernel 的最佳化可以分成四個方向：weight packing 調整資料配置，L1 tiling 和 prefetch 處理記憶體存取，RVV 在單一 CPU core 內做 SIMD，at::parallel_for 再把工作分給多個 core。

at::parallel_for：把 MB × NB tiles 分給多個 CPU core
        -> RVV intrinsic：每個 core 同時計算多個輸出通道
        -> K-loop unroll / prefetch：減少 pipeline 等待資料的時間
        -> packing / tiling：讓 weight 連續載入並改善 cache locality

Vector-256bit 與 RVV 分塊

K1 的 VLEN 是 256 bit。這份實作將 BLOCK_N 定義成 __riscv_v_fixed_vlen / 4，所以在 K1 上是 64 個輸出通道。通用 GEMM kernel 使用 FP32 LMUL=4 累加器，一個 block 分兩次向量迴圈處理；decode 常見的 M=1 另有 LMUL=8 的 GEMV 路徑，可以一次覆蓋 64 個 FP32 元素。

BF16 / FP16 weight 會先載入較窄的向量暫存器，轉成 FP32 後再以 vfmacc_vf 累加。轉成 FP32 可以保留較好的累加精度，也會增加 vector register 壓力，所以 LMUL 和 unroll 不能無限放大。

用 L1 cache size 決定 K tile

K1 每個 X60 core 有 32KB L1-D cache。這份 CMake 也將 RVV_L1_CACHE_KB 預設為 32，並允許在編譯時覆寫。gemm.cpp 再用它計算 K 方向的 tile size KB。

以 VLEN=256、BLOCK_N=64 的 BF16 路徑為例：

KB = 32KB × 2 / 256 = 256
packed-weight tile = KB × BLOCK_N × 2 bytes
                   = 256 × 64 × 2
                   = 32KB

我目前設計是讓 packed-weight tile 約等於 L1-D cache 大小。這是編譯時的 heuristic，activation、C_f32 與其他資料也會佔用 cache，因此實際 cache hit rate 還是要透過 benchmark 測量。

Prefetch、unroll 與 in-order pipeline

X60 是雙發射、順序執行（in-order）的 CPU core。遇到相依的 load 與 FMA 時，pipeline 可能需要等待資料。這份 kernel 會把 K-loop 展開成兩組或四組獨立的 weight load 與 vfmacc，讓 compiler 有更多可排程的指令。

常見的 Intel Core 或 Xeon 使用亂序執行（out-of-order execution），CPU 可以在執行時尋找其他已準備好的指令，先填補等待 load 的空檔。X60 的 in-order core 主要依靠 compiler 在編譯時排好指令順序，因此 kernel 需要用 unroll 主動提供多組彼此獨立的 load 和 FMA。如果指令之間的相依太緊密，in-order pipeline 會更容易停頓。

迴圈內也有 software prefetch

if (k + PREFETCH_DIST < K) {
  __builtin_prefetch(
      b_ptr_base + (k + PREFETCH_DIST) * BLOCK_N, 0, 1);
}

它會在計算目前 K 位置時，提前請求後面的 packed weight。Prefetch 和 unroll 的目標是降低等待記憶體的時間，不會減少實際需要讀取的 weight 量。這條編譯路徑使用 -O3 -march=rv64gcv，目前沒有 -mcpu=spacemit-x60 或 -mtune=spacemit-x60，所以可以說 kernel 的寫法有考慮 in-order pipeline，compiler 還沒有啟用 X60-specific instruction scheduling。

OpenMP 與 RVV 在不同範圍平行化

OpenMP 與 RVV 處理不同範圍的平行工作。

gemm.cpp 沒有直接寫 #pragma omp parallel for，它使用 PyTorch 的 at::parallel_for。CMake 有加入 -fopenmp 並連結 OpenMP；當目標 PyTorch 使用 OpenMP parallel backend 時，at::parallel_for 才會由 OpenMP 執行。因此一個執行緒處理一組 tile 時，還可以在同一個 core 內使用 RVV。

多核心不保證線性加速。不同 core 會共用 L2 cache 與 LPDDR 頻寬，像 gate/up projection 和 lm_head 都需要串流讀取很大的 weight。執行緒數增加後，最後可能卡在共用記憶體頻寬。

lm_head 還有一個資料配置限制。很多模型讓 lm_head.weight 與 embedding 共享同一個 tensor，不能在載入時直接把共享 weight 原地替換成 packed layout。目前實作對 tied lm_head 使用懶建立的 side cache，並要確保 torch.compile 重播時不會重複 packing。

效能測試要怎麼看

看 linear 效能測試時，我會分成三層。

第一層是 kernel 微型效能測試。它能回答單一形狀下 RVV kernel 是否比基準快。

第二層是 運算子層級效能分析。它能回答 SGLang 前向傳播 裡是否真的走到 sglang.rvv.linear.weight_packed_linear。

第三層是 單批次 / 推論服務效能測試。它能回答整個推論服務路徑是否變快。

這三層要一起看。微型效能測試變快只代表 kernel 本體有希望；如果載入時 packing、執行環境分派、cache 壓力、lm_head、attention 或計算圖擷取成本沒有處理，端到端結果仍然可能看不到同等比例的改善。

今天先走到這裡

今天把 SGLang RVV linear 路徑接起來

• LLM 裡的 QKV projection、MLP projection、lm_head 都會大量使用 linear。
• SGLang RVV 後端會在載入時用 convert_weight_packed 預處理 weight，結果可能是 blocked-packed 或 FP32-transposed representation。
• Python 前向傳播裡透過 torch.ops.sgl_kernel.weight_packed_linear 呼叫 C++ RVV kernel。
• C++ kernel 會檢查形狀、資料型別與預處理 representation，再分派到 FP32 / BF16 / FP16 的 blocked-packed 或 FMA fallback RVV 路徑。
• 這份 kernel 用 weight packing、L1-aware K tiling、prefetch、K-loop unroll 和 RVV 改善單核心計算，再用 at::parallel_for 切分多核心工作。
• OpenMP 負責 thread-level 平行，RVV 負責單核心內的 data-level SIMD；核心數增加後還要注意共用 L2 cache 和 LPDDR 頻寬。

明天看 norm、activation、RoPE。那些運算子的計算量比 GEMM 小，但它們出現頻率高，而且會碰到 reduction、近似函式和資料配置，寫法和 linear 很不一樣。

參考資料

• SGLang PR #22193: Add RVV kernels in sgl-kernel
• Banana Pi BPI-F3 SpacemiT K1 Datasheet
  參考 Repo
• python/sglang/srt/layers/rvv_utils.py
• python/sglang/srt/layers/quantization/unquant.py
• sgl-kernel/csrc/cpu/riscv64/gemm.cpp
• sgl-kernel/csrc/cpu/riscv64/vector_helpers.h
• sgl-kernel/csrc/cpu/riscv64/CMakeLists.txt