iT邦幫忙

2026 iThome 鐵人賽

DAY 10
0
Software Development

在 AI Compiler 工程師的路上系列 第 10

Day09:把 TorchInductor RVV 接進 SGLang:Packed State、Small Batch 和重現環境

  • 分享至 

  • xImage
  •  

昨天完成 P1–P3,讓 TorchInductor 能選到 VecRVV,BF16 Linear 有 M=1M>=2 microkernel,靜態 shape gate 也會保留不支援情況的 fallback。

今天往 runtime 繼續,microkernel 能跑之後,還有三個部署問題

P4  weight 要轉成 blocked layout,誰負責持有與失效?
P5  如何讓使用者在原本的 SGLang TorchNative path 用一個 flag 就能處理細節?
P6  換一個乾淨環境後,原始碼、image、model 和 cache 定義能不能重現?

可以參考的 branch

本篇大綱

  • 先看 row-major weight 為什麼限制 RVV Linear
  • 讀 P4 的 explicit packed tensor 與 SGLang module ownership
  • 檢查 reload、mutation、tied weight 和記憶體預算
  • 讀 P5 如何用一個 opt-in 安裝 SGLang policy,再看 2/4/8/16 bucket、padding / slicing 與 lm_head routing
  • 用 one-batch 和 serving 結果分開回答 latency、throughput 與 RSS
  • 最後用 P6 的固定 source、image、checkpoint、cache state 重現一次

P4:Runtime weight 為什麼需要明確 owner

一般 Linear weight 是 row-major [N, K]。昨天的 RVV BF16 template 最適合沿 N 連續載入一段 weight vector,所以 P4 將 weight 轉成適合向量化的 blocked layout。Packed-weight primitive 與 layout contract 的完整實作可以參考 P4 commit

row-major: [N, K]
packed:    [ceil(N / 32), K, 32]

最後一維的 32 個輸出 channel 連續存放,剛好對應 microkernel 的 block_n=32。PyTorch 的 pack primitive 如下

RVV_BF16_PACKED_WEIGHT_BLOCK_N = 32
RVV_BF16_PACKED_WEIGHT_LAYOUT_VERSION = 1


def eager_rvv_pack_bf16_weight(weight, layout_version=1):
    n, k = weight.shape
    block_n = RVV_BF16_PACKED_WEIGHT_BLOCK_N
    padded_n = (n + block_n - 1) // block_n * block_n
    padded_weight = torch.nn.functional.pad(
        weight, (0, 0, 0, padded_n - n)
    )
    return (
        padded_weight
        .reshape(padded_n // block_n, block_n, k)
        .permute(0, 2, 1)
        .contiguous()
    )

假設 W.shape == [3072, 2048],packed tensor 會是

[3072 / 32, 2048, 32] = [96, 2048, 32]

如果 N 不能被 32 整除,先把 N pad 到下一個 32。呼叫 Linear 時仍傳入原本的 out_features,輸出會忽略 padding 部分。

PyTorch 將 pack 與 packed Linear 做成兩個 Inductor prim,做成 Inductor 可以認得的自訂基本操作。

rvv_pack_bf16_weight = make_prim(
    "rvv_pack_bf16_weight(Tensor weight, int layout_version=1) -> Tensor",
    eager_rvv_pack_bf16_weight,
)

rvv_packed_bf16_linear = make_prim(
    "rvv_packed_bf16_linear("
    "Tensor input, Tensor packed_weight, "
    "SymInt out_features, int layout_version=1) -> Tensor",
    eager_rvv_packed_bf16_linear,
)

layout version 是 runtime contract,這份 packed tensor 屬於 runtime layout,不應寫進 checkpoint,model load 完成後,要由使用它的 runtime 重新建立。

Module-owned side state

這裡早期設計是把 cache 放在 generated C++ function-local static map,key 依賴 raw pointer 和 fingerprint。這個設計很難完整掌握 Tensor lifetime、reload 與 mutation,會發生問題像是模型刪掉,但 static map 還是存在,packed weight 可能無法及時釋放。載入新模型後,allocator 可能重複使用舊記憶體位址,cache 可能誤認成同一個 weight。

新版把 packed tensor 交給 SGLang module 持有,好處是 module 存在,packed weight 就存在,module 被釋放,packed weight 也跟著釋放,而且 PyTorch 能透過 module 管理 buffer,reload 時可以用 module hook 明確標記 packed state 失效。執行前可以檢查目前 module.weight 是否仍符合當初的 pointer、shape、stride、dtype 和 device。

def _register_explicit_packed_weight(module):
    weight = module.weight
    with torch.inference_mode(False), torch.no_grad():
        packed_weight = _RVV_PACK_BF16_WEIGHT_OP(weight.detach())

    module.register_buffer(
        "_sglang_rvv_packed_weight",
        packed_weight,
        persistent=False,
    )
    module._sglang_rvv_packed_source = _PackedWeightSource.capture(weight)

這裡有三個 ownership 決定。

1.packed weight 以 tensor 表示,沒有藏在 C++ static map 裡。它跟著 module 移動與釋放,PyTorch allocator 也看得到這份記憶體。

2.buffer 設為 persistent=Falsestate_dict() 不會把 runtime-only layout 寫進 checkpoint,避免 layout version、target ISA 或 padding 規則滲進模型格式。

3.原始 row-major module.weight 仍然保留。packed state 失效或記憶體預算不足時,可以退回 F.linear(input, weight)

但這有代價的,同一組 weight 同時保留 row-major 與 packed layout,peak RSS 會增加。後面的 Llama 實驗量到約 +2.30 GiB,保留兩份是為了相容性,PyTorch fallback或其他算子使用,或是 tied weight 共用像是 embedding 和 lm_head。

怎麼判斷 packed state 已經過期

SGLang 會記錄 packed tensor 的來源

@dataclass(frozen=True)
class _PackedWeightSource:
    tensor_ref: weakref.ReferenceType
    data_ptr: int
    shape: tuple[int, ...]
    stride: tuple[int, ...]
    dtype: torch.dtype
    device: torch.device

    def matches(self, weight):
        return (
            self.tensor_ref() is weight
            and self.data_ptr == weight.untyped_storage().data_ptr()
            and self.shape == tuple(weight.shape)
            and self.stride == tuple(weight.stride())
            and self.dtype == weight.dtype
            and self.device == weight.device
        )

實際呼叫前會檢查

def _valid_explicit_packed_weight(module):
    weight = module.weight
    packed_weight = module._buffers.get("_sglang_rvv_packed_weight")
    source = getattr(module, "_sglang_rvv_packed_source", None)

    if packed_weight is None:
        return None
    if not isinstance(source, _PackedWeightSource):
        return None
    if not source.matches(weight):
        return None
    return packed_weight

如果 weight tensor 被替換、storage 改變、shape / stride 改變、dtype 或 device 改變,函式回傳 None,regional policy 便使用 row-major fallback。

model reload 另有明確 invalidation hook

def invalidate_rvv_inductor_packed_weights(model, *_):
    for module in model.modules():
        if "_sglang_rvv_packed_weight" in module._buffers:
            module._sglang_rvv_packed_source = None


model.register_load_state_dict_post_hook(
    invalidate_rvv_inductor_packed_weights
)

原地修改 weight 時,tensor identity 和 storage pointer 可能都不變,所以單靠 _PackedWeightSource.matches() 無法偵測內容變化。支援的公開 weight lifecycle 必須呼叫 invalidation,目前是假設模型完成載入後,推論期間 weight 不會被任意原地修改。對一般唯讀推論這個假設通常成立,未來如果要成為更完整的設計,所有 weight mutation、reload、量化或 adapter 更新都必須觸發 invalidation 或重新 pack。

Tied lm_head 為什麼需要另外檢查

部分模型設定 tie_word_embeddings=True,要去看具體 checkpoint 的 config.json 才知道,embed_tokenslm_head 共用同一份 weight。安裝 policy 時不能偷偷換掉 lm_head.weight,否則 weight tying 會被破壞。

SGLang 會先驗證 module alias

def _validate_tied_lm_head(model):
    if not model.config.tie_word_embeddings:
        return
    embed_tokens = model.model.embed_tokens
    if model.lm_head is not embed_tokens:
        raise RuntimeError(
            "tied lm_head and embed_tokens must remain the same module"
        )

packed tensor只是 side buffer,原始 weight object 和 tied module identity 都保持不變。lm_head 可以用自己的 packed side state做投影,embedding lookup 仍使用原始 weight。

記憶體不足時只 pack 一部分

explicit packed state 會增加 RSS(Resident Set Size),RSS是指說一個 Process 實際佔用並停留在實體記憶體中的大小,所以 P4 先計算每個 module 所需空間

padded_out_features = (
    (out_features + block_n - 1) // block_n * block_n
)
packed_bytes = padded_out_features * in_features * weight.element_size()

使用者可以設定:

--cpu-rvv-memory-budget-mib N

沒有明確上限時,SGLang 讀取目前還可以用多少的主機記憶體,超過預算的 projection 不建立 packed tensor,執行時走 row-major。這可以涵蓋部分形狀,不需要因為一個 module 放不下就讓整個 model 啟動失敗,不過目前還沒有測試部分 pack ,會影響多少性能。

P5:用原本 TorchNative 參數接上 Inductor RVV

P5 把 P1–P4 接到 SGLang 原本的 TorchNative model-load 與 serving path。完整整合可以參考 P5 commit。這個

使用者保留原本的 CPU、BF16 與 TorchNative attention 參數,只加一個 flag

sglang serve \
  --model-path "${MODEL}" \
  --device cpu \
  --dtype bfloat16 \
  --attention-backend torch_native \
  --enable-cpu-rvv-inductor

--enable-cpu-rvv-inductor 會選擇已測試的 regional compile、explicit packed BF16 weights 與 M=2–15 buckets。使用者不需要分別理解或設定 --cpu-compile-mode regional--cpu-rvv-packed-weight-mode explicit 和 small-batch option。

Policy 在正常 checkpoint load 完成後安裝。它解析 model adapter 和 projection modules,建立 module-owned packed side state,檢查 memory budget 和 tied lm_head,再依 shape key 建立或重用 compiled callable。條件不符合或 packed state 無效時,該 projection 回到 row-major F.linear

兩條 decode sequence 為什麼形成 M=2

單 request decode 的 hidden state 通常是

[1, hidden_size]

當 scheduler 同時處理兩條 decode sequence,Linear 看到的 flattened rows 會變成

[2, hidden_size]

舊 regional policy 主要處理 M=1 decode 和 M=16–128 prefill。M=2–15 會退回 eager,導致小批次 serving 沒有用到上一篇的 M>=2 microkernel。

P5 將 row count 映射到有限 bucket

def _regional_bucket_rows(rows, *, allow_small_batch=False):
    if rows == 1:
        return 1

    buckets = (
        (2, 4, 8, 16, 32, 64, 128)
        if allow_small_batch
        else (16, 32, 64, 128)
    )
    for bucket in buckets:
        if rows <= bucket:
            return bucket
    return rows

實際映射如下

實際 M 編譯 bucket
1 1
2 2
3–4 4
5–8 8
9–15 16
16 16
17–32 32
33–64 64
65–128 128

M=3 不會產生一份專屬 artifact。runtime 先 pad 到 4,執行完成後切回前 3 列

m, n, k = _linear_shape(input, weight)
bucket_m = _regional_bucket_rows(
    m, allow_small_batch=self.allow_small_batch
)

compile_input = input
if bucket_m != m:
    compile_input = F.pad(input, (0, 0, 0, bucket_m - m))

output = compiled(compile_input, packed_weight, n)

if bucket_m != m:
    output = output[:m]

這是一個以額外運算換 artifact 重用的選擇。M=9 會用 16-row kernel,多算 7 列;好處是 M=9–15 共用同一份 compiled callable。

Regional shape key 和 fallback

compiled callable 依這組 key 重用

class _RegionalShapeKey(NamedTuple):
    dtype: torch.dtype
    device_type: str
    projection: str
    rows: int
    out_features: int
    in_features: int
    mode: str

mode 會區分 row_majorexplicit_packed,避免兩種 weight contract 共用錯誤 artifact。projection 則限制在

("qkv_proj", "o_proj", "gate_up_proj", "down_proj", "lm_head")

shape gate 只接受

CPU
BF16 input / BF16 weight
2D input / 2D weight
bias is None
N >= 1024
K >= 1024
M × N × K >= 1,000,000
M = 1、opt-in M = 2–15,或非 lm_head 的 M = 16–128

其他 dtype、device、shape、bias 或 projection 都走 F.linear。P5 特別把 lm_head 加進 M=2–15,如果只處理 transformer block 裡的四種 projection,兩條 sequence 最後仍會在詞彙投影退回 eager。

前面使用者打開的 flag

--enable-cpu-rvv-inductor

會一次啟用 regional compile、small-batch buckets 與 explicit packed BF16 weights,完整 CPU graph 維持關閉。

One-batch:Packed layout 帶來多少差異

Llama-3.2-1B 的 final row-major 與 explicit-packed rows 使用相同 source、wheel、checkpoint、workload 和 cache 定義

實驗設定

batch = 1
input_len = 64
output_len = 8
BF16, TP=1
torch_native attention
OMP_NUM_THREADS=8
regional compile
full CPU graph off

https://ithelp.ithome.com.tw/upload/images/20260810/20183319emD8r3FZhE.png

在這個 Llama one-batch 條件下,明確持有 blocked weight 可以改善 Inductor RVV Linear,代價是多保留一份 weight layout。

Qwen3-0.6B 和 DeepSeek-R1-Distill-Qwen-1.5B 也完成三個 artifact-warm fresh process

Model Process wall median Prefill median Decode median Total median
Qwen3-0.6B 241.12 s 13.513 tok/s 1.016 tok/s 6.195 tok/s
DeepSeek-R1-Distill-Qwen-1.5B 272.30 s 8.999 tok/s 0.871 tok/s 4.710 tok/s

這兩列證明路徑不只套在一個 model class。

Serving:Concurrency 2 改善哪一個指標

Llama serving 測試讓每個新開啟的 server 依序跑 concurrency 1 和 2,每個 concurrency 有四個 measured requests,每個 request 固定 64 input、8 output tokens。

Metric c1 median c2 median c2/c1
Benchmark wall 49.114 s 37.785 s -23.1%
TTFT 6.606 s 12.125 s +83.5%
Per-request decode 1.255 tok/s 1.060 tok/s -15.5%
Global decode 1.258 tok/s 1.979 tok/s +57.4%
Aggregate total 5.864 tok/s 7.622 tok/s +30.0%

Concurrency 2 提高 global decode 和 aggregate total throughput,也縮短整批 benchmark wall。每個 request 的 TTFT 與 decode rate 反而變差,因為兩個 request 共用 CPU 資源。

還有發現系統總吞吐提高,每個 request 的等待與 decode latency 變差。bucket routing 讓 M=2 有 RVV kernel 可用;是否接受較高 TTFT,取決於服務目標偏 throughput 或 per-request latency。

P6:固定 source、image、model 和 cache state

P6 整理 Banana Pi 上的 clean user reproduction path,包含環境、image、model 與執行入口。完整變更可以參考 P6 commit

如果想自己測試,可以準備一台 Banana Pi BPI-F3(16GB RAM)。實驗還需要 Git、Podman,以及能下載模型與 container image 的網路環境。

重現時要固定四個條件:PyTorch / SGLang source branch、riscv64 runtime image、model checkpoint revision,以及 Inductor cache 是 artifact-cold 還是 artifact-warm。cold run 從空 cache 開始,warm run 則在新的 container / process 裡重用同一份 cache。

完整的實驗重現教學放在 README_TORCH_INDUCTOR_RVV.md。README 已整理好環境檢查、source clone、固定 image、Qwen3-0.6B one-batch、generated artifact 驗證,以及 serving concurrency 1 / 2 的操作步驟。

P1–P6 回顧

https://ithelp.ithome.com.tw/upload/images/20260810/20183319nrQbLQD4A8.png

P1–P4 的主要修改位於 PyTorch/TorchInductor,P5 把這些能力和 SGLang 原本的 TorchNative model-load 與 serving flow 整合, P6 再固定可重現的使用環境。

Phase 主要 repo 解決的缺口 使用者可觀察結果
P1 — VecRVV / CPU ISA PyTorch TorchInductor 不會選 RVV,也不會帶入對應 macros 與 -march torch.compile 產生的 CPU C++ 能選到 RVV target
P2 — BF16 microkernels PyTorch 一般 ISA 支援沒有 LLM Linear 的計算核心 Decode M=1 使用 GEMV;prefill/small batch M>=2 使用 GEMM
P3 — Static template routing PyTorch 所有 shape 都強迫編譯會造成錯誤選路與 artifact 膨脹 支援的 static BF16 Linear 進 RVV template;其餘保留 fallback
P4 — Explicit packed state contract PyTorch + SGLang boundary Runtime row-major weight 重複付出 layout conversion,packed data 缺少 owner 與失效規則 PyTorch 提供 layout v1 primitives;SGLang module 持有 non-persistent side state
P5 — SGLang native integration policy SGLang Compiler primitives 無法直接從原本 serving CLI 使用 保留 --device cpu --dtype bfloat16 --attention-backend torch_native,只加 --enable-cpu-rvv-inductor;policy 管理 adapter、budget、lm_head、small batch 與 fallback
P6 — Clean reproduction SGLang / environment 更容易重現實驗結果 固定 public branch、image digest、model revision、cold/warm cache 與重現命令

今天先走到這裡

今天完成 P4–P6

  • P4 把 BF16 weight pack 成 [ceil(N/32), K, 32],由 SGLang module 持有 non-persistent side buffer
  • packed source 會檢查 tensor identity、storage、shape、stride、dtype 和 device;reload 或公開 weight lifecycle 會使它失效並退回 row-major
  • tied embed_tokens / lm_head 保持原本 module 和 weight identity,packed tensor 只作為 side state
  • P5 保留原本 TorchNative model-load 與 serving API;使用者只加 --enable-cpu-rvv-inductor,policy 負責 adapter、packed lifecycle、memory budget、routing、lm_head 與 fallback
  • P5 的 small-batch 子路徑把 M=2–15 映射到 2/4/8/16,用 padding / slicing 限制 compiled artifact 數量
  • Llama one-batch 顯示 explicit packed 改善 prefill、decode 和 total throughput,同時增加約 2.30 GiB peak RSS
  • concurrency 2 提高 global throughput,但 TTFT 和 per-request decode 變差

明天會回顧這九天所學,再複習一下有趣的知識。

參考資料


上一篇
Day08:PyTorch Conference 2026 Poster:為什麼要讓 TorchInductor 支援 RVV
下一篇
Day10:SGLang RVV 和 TorchInductor RVV 複習日
系列文
在 AI Compiler 工程師的路上19
圖片
  熱門推薦
圖片
{{ item.channelVendor }} | {{ item.webinarstarted }} |
{{ formatDate(item.duration) }}
直播中

尚未有邦友留言

立即登入留言