前兩天先學硬體和基本 RVV intrinsic，今天開始看要怎麼接進真實推論服務系統 SGLang。
今天先看懂一個 RVV attention 後端在 SGLang 裡要接到哪些位置，接下來幾天再看怎麼用 RVV Intrinsic 手寫 GEMM、norm、activation、RoPE、attention 等等 kernel，了解為什麼可以變成 SGLang 可以呼叫的 torch.ops.sgl_kernel.*。

本篇大綱

• RVV attention backend 支援接進 SGLang
• SGLang 的前端、後端執行環境、RadixTree 與 RadixAttention，以及 attention 後端在推論服務裡負責哪一段。
• 為什麼 RISC-V CPU 需要 RVV attention 後端。
• Python 到 C++ 的流程，看後端註冊、Python 分派、PyTorch 自訂運算子綁定、RVV kernel 。
• 最後整理這個後端怎麼和 SGLang 整合，以及接下來幾篇會拆哪些 AI 運算子。

RVV attention 後端接進 SGLang

這幾天的內容以 kernel 工程為主，先了解手寫 kernel，後續介紹 AI compiler 自動產生程式碼時，就能比較兩種實作方式。
我會介紹我目前正在 SGLang 的 issue 和 PR 中推進 RVV attention 後端，來解釋這個後端如何整合，這些內容目前還沒有進入 SGLang 正式版還在努力中。

• Issue：sgl-project/sglang#18072
• Phase 1 PR：sgl-project/sglang#22193
• Phase 2 PR：sgl-project/sglang#22219

我在 Issue #18072 有提到我把工作拆成幾段：先做基礎 RVV kernel 和單元測試，再把這些 kernel 接進 SGLang 的 Python 分派層，最後補 Dockerfile 和 document。

今天先看整合面。接下來幾篇再拆運算子，會比較容易看出每個 intrinsic 最佳化技巧要放在哪一層。

SGLang 是什麼?

SGLang 是一個專為 LLM 和多模態模型設計的高效能推論與服務框架，SGLang 的名稱來自 Structured Generation Language。它提供模型服務 API，也能描述、排程和執行語言模型程式。語言模型程式可以包含多輪生成、分支、平行呼叫、選擇候選答案、結構化輸出，以及把生成結果接到下一輪輸入。

詳細可以看論文〈SGLang: Efficient Execution of Structured Language Model Programs〉，有提到 SGLang 被拆成兩個互相配合的部分：前端語言與後端執行環境，共同思考如何一起最佳化。

前端:描述要怎麼生成

SGLang 的前端是一套嵌入 Python 的領域專用語言。使用者可以用 gen、select、fork、join 等 primitive，描述生成、選擇與平行控制流程，同時保留 Python 原本的控制流程與函式庫。

例如一個工作流程可能先產生三個候選答案，再平行檢查每個答案，最後選出結果。從前端來看，我們關心的是 prompt 怎麼組合、生成流程怎麼分支，以及不同模型呼叫之間有哪些相依關係。

現在許多使用者也會直接透過 OpenAI-compatible API 使用 SGLang，把它當成模型 serving framework。無論請求來自 SGLang 前端程式或 API，最後都需要交給後端執行環境完成排程、模型前向傳播和 token 生成。

後端：決定怎麼有效率地跑

後端執行環境接到請求後，會處理 batching、排程、KV cache 管理、模型執行、結構化輸出解碼，以及不同硬體 kernel 的分派。前端描述的是應用程式想做什麼；後端則要把這些生成操作轉成實際的 prefill、decode 和模型運算。

整體執行路徑可以整理成

SGLang 前端程式 / OpenAI-compatible API
-> SGLang Runtime：排程、batching、KV cache 管理
-> Model Runner：執行模型每一層
-> Attention Backend：選擇 decode / extend 實作
-> CUDA / Triton / Torch Native / RVV kernel

這篇要接進來的 RVV attention backend 位於較下層的執行層。runtime 先整理 query、KV cache 和請求 metadata，再交給 RVV attention backend 完成 RISC-V CPU 上的 attention 計算。

RadixTree 與 RadixAttention

SGLang 最有名的 runtime 最佳化之一是 RadixAttention。它負責管理 不同請求之間的 KV cache 重用，attention 的數學公式維持不變。

LLM 在處理 prompt 時，會為每個 token 計算 attention 所需的 key 和 value。只要兩個請求有相同的 token 前綴，這一段前綴對應的 KV cache 就可以重用，不必重新做一次 prefill。例如多位使用者可能共用同一段 system prompt；多輪對話也會重複帶入先前的聊天紀錄。

SGLang 使用 RadixTree 管理 token 序列與 KV cache 的對應關係。一般 trie 的一條 edge 常只放一個元素，RadixTree 的一條 edge 可以保存一段 token 序列，因此能更緊湊地表示大量共享前綴。

[共同 system prompt]
├── [使用者問題 A] -> KV cache A
└── [使用者問題 B] -> KV cache B

新請求進來時，runtime 先在 RadixTree 裡做最長前綴匹配。已經存在的前綴可以直接沿用 KV cache，只需要對後面新增的 token 做計算。請求完成後，cache 會留在樹中供後續請求重用；空間不足時再配合 LRU eviction 移除較久未使用的節點。runtime 也可以優先排程共享前綴較長的請求，提高 cache hit rate。

RadixTree、RadixAttention 和 attention backend 分別負責以下工作：

• RadixTree：保存 token 前綴與 KV cache 對應關係的資料結構。
• RadixAttention：利用 RadixTree、cache-aware scheduling 與 eviction policy，自動進行 KV cache 重用的 runtime 機制。
• Attention backend：拿到 query 和已配置好的 KV cache 後，執行 attention 計算的硬體實作。

RadixAttention 在 runtime 層決定哪些前綴可以重用，以及 token 對應到哪一段 KV cache。RVV backend 在較下層讀取這些資料，執行 decode 或 extend attention kernel。

今天會專注在後端執行環境的一小段：模型前向傳播時，SGLang 怎麼選擇 attention backend，並一路呼叫到 C++ RVV kernel。

為什麼要做 RVV attention 後端

上面是示意圖，SGLang 原本支援很多後端，例如 FlashInfer、Triton、Intel AMX、Ascend、XPU。這些後端大多針對 GPU、特定加速器，或 x86 CPU 的向量/矩陣指令。

RISC-V CPU 上沒有 CUDA，也不能直接使用 GPU Triton kernel 。像 Banana Pi BPI-F3 / SpacemiT K1 這類板子有 RVV 1.0 和 256-bit 向量，對推論服務來說需要一條 CPU 原生路徑

Python 模型程式碼
-> SGLang 執行環境
-> RISC-V CPU 後端
-> C++ RVV intrinsic kernel 
-> RVV instruction

我的目的是希望可以加速在 RISC-V CPU 上用 SGLang 執行模型，原本只有純 pytorch 寫的後端可以跑，但對於 RISC-V CPU 來說沒有利用到 RVV 加速有點可惜。

最佳化後可以把重複出現的運算整合成 C++ 自訂運算子，並在內部使用 RVV 與特定資料配置。效能改善可能來自運算融合、減少中間 tensor、較適合 cache 的配置與 RVV 向量化。單純將程式從 Python 移到 C++ 不保證降低成本，還是要用 profiler 與端到端測試驗證。

kernel 也可以根據 RISC-V CPU 的特性調整。K1 的 Vector-256bit、L1/L2 cache、LPDDR 頻寬、dual-cluster 都會影響分塊、執行緒切分、載入/儲存 型態和量化資料處理。

SGLang 可以保留原本的高階推論服務流程。使用者仍然透過 python3 -m sglang.launch_server 啟動服務，只是在 --attention-backend rvv 時，執行環境會把 attention 送到 RVV 後端。

今天要讀的檔案

接下來幾篇會再進到 gemm.cpp、norm.cpp、activation.cpp、rope.cpp 和 attention kernel 內部。

SGLang 怎麼找到 RVV 後端

SGLang 的 attention 後端透過註冊表管理。RVV 後端會被註冊成名稱 rvv，之後使用者可以透過啟動參數選它：

python3 -m sglang.launch_server \
    --model-path Qwen/Qwen2.5-1.5B-Instruct \
    --attention-backend rvv \
    --device cpu \
    --dtype bfloat16

在程式碼裡，註冊表是長這樣：

@register_attention_backend("rvv")
def create_rvv_backend(runner):
    from sglang.srt.layers.attention.rvv_backend import RVVAttnBackend

    return RVVAttnBackend(runner)

SGLang 執行環境不需要知道 RVV kernel 的細節。它只要拿到一個符合 AttentionBackend 介面的物件。後面 decode / extend 的時候，執行環境會呼叫這個後端的 forward_decode 或 forward_extend。

RVVAttnBackend 初始化做了什麼

RVVAttnBackend 在初始化時先設定幾個狀態：

• device = "cpu"
• use_rvv_kernels = False
• Fallback 後端使用 TorchNativeAttnBackend
• 如果 cpu_has_rvv_support() 通過，就嘗試初始化 RVV kernel

初始化時會去拿 torch.ops.sgl_kernel 裡的運算子

self.decode_fwd_impl = ops.decode_attention_cpu
self.extend_fwd_impl = ops.extend_attention_cpu

如果 KV cache是 INT8，會改拿

self.decode_fwd_impl = ops.decode_attention_int8_cpu
self.extend_fwd_impl = ops.extend_attention_int8_cpu

這個設計讓後端可以在 kernel 缺少時退回 Torch Native，因為不同機器可能沒有編好 sgl-kernel，或硬體沒有 RVV。程式要能把錯誤說清楚，不能只在推論服務啟動後直接炸掉。

初始化還會建立 _attn_logits_pool。這是一個預先配置的 CPU tensor，用來存 decode kernel 的 部分 softmax buffer。它也服務後面的 CPU 計算圖 / torch.compile 路徑，因為計算圖重播需要穩定的tensor形狀和資料指標。

Python 執行 decode 的流程

decode 是LLM推論服務裡很常見的狀態：每個請求已經有一段 KV cache，現在要產生下一個 token。

在 forward_decode 裡，RVV 後端會做幾件事

1. 檢查 RVV kernel 是否啟用
2. 決定這次 K/V 要寫到哪個 KV cache 位置
3. 把 q reshape 成 [batch, head, head_dim]
4. 建立輸出 tensor
5. 從 forward_metadata 拿到 attn_logits buffer
6. 組好 C++ 運算子需要的參數
7. 呼叫 self.decode_fwd_impl(*args)

可以把它畫成這樣：

forward_decode(q, k, v, 層, forward_batch)
-> q.view(...)
-> 輸出 tensor
-> token_to_kv_pool.get_key_buffer(layer_id)
-> token_to_kv_pool.get_value_buffer(layer_id)
-> req_to_token / req_pool_indices / seq_lens
-> torch.ops.sgl_kernel.decode_attention_cpu(...)

這裡的 req_to_token、req_pool_indices、seq_lens 很重要。attention kernel 需要知道每個請求對應到哪些token 位置，才能從 KV cache裡找出正確的 key/value。

Python 執行 extend 的流程

extend 對應 prefill 或 chunked prefill 類情境，它一次處理多個新 token，會碰到前綴 KV cache和本輪 extendtoken。

forward_extend 的流程和 decode 類似，但參數多了幾個和 extend 長度有關的 tensor：

forward_extend(q, k, v, layer, forward_batch)
-> q.view(...)
-> output tensor
-> k_extend / v_extend
-> k_buffer / v_buffer
-> extend_seq_lens
-> extend_start_loc
-> max_extend_len
-> torch.ops.sgl_kernel.extend_attention_cpu(...)

extend kernel 要處理兩段資料：

• 前綴：已經在 KV cache 裡的token。
• 目前 extend：這一輪新進來的一段 token。

所以它比 decode 更像一個小型 tiled attention。後面的 attention 實作篇會進去看 extend.cpp 裡的分塊 buffer、softmax update、BLOCK_M / BLOCK_N 設計。

C++ 怎麼暴露成 torch.ops

Python 端呼叫的 torch.ops.sgl_kernel.decode_attention_cpu 來自 PyTorch 自訂運算子註冊，實作在 C++ 擴充 裡。

在 sgl-kernel/csrc/cpu/riscv64/torch_extension_riscv64.cpp 裡，可以看到這種註冊：

TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
  m.def(
      "decode_attention_cpu(Tensor query, Tensor(a!) k_cache, Tensor(b!) v_cache, "
      "Tensor(c!) output, Tensor key, Tensor value, Tensor loc, "
      "Tensor(d!) attn_logits, Tensor req_to_token, Tensor req_pool_indices, "
      "Tensor seq_lens, float sm_scale, float logit_cap) -> ()");
  m.impl("decode_attention_cpu", torch::kCPU, &decode_attention_cpu);

  m.def(
      "extend_attention_cpu(Tensor q_extend, Tensor k_extend, Tensor v_extend, "
      "Tensor(a!) o_extend, Tensor k_buffer, Tensor v_buffer, "
      "Tensor req_to_token, Tensor req_pool_indices, Tensor seq_lens, "
      "Tensor extend_seq_lens, Tensor extend_start_loc, int max_len_extend, "
      "float sm_scale, float logit_cap) -> ()");
  m.impl("extend_attention_cpu", torch::kCPU, &extend_attention_cpu);
}

m.def 定義 operator schema，包含參數、回傳值與 alias / mutation 註記；m.impl 再把這個 schema 接到 CPU 的 C++ 實作。Tensor(a!) 這種標註表示該 tensor 會被 in-place 寫入。這會影響 PyTorch 別名分析，也會影響計算圖 capture 和compiler對副作用的理解。

這一步接好之後，Python 端就可以用：

torch.ops.sgl_kernel.decode_attention_cpu(...)

呼叫 C++ RVV kernel

C++ kernel 裡開始出現 RVV intrinsic

進到 decode.cpp 或 extend.cpp 後，才會看到前兩篇講過的 RVV intrinsic。

例如 decode attention 的 tiny GEMM 會沿著 head 維度 做 strip-mining：

size_t vl;
for (int64_t k = 0; k < K; k += vl) {
  vl = __riscv_vsetvl_e32m1(K - k);
  ...
  vc[m * COLS + n] =
      __riscv_vfmacc_vv_f32m1_tu(vc[m * COLS + n], va, vb[n], vl);
}

這段和昨天的 add_f32 迴圈是同一個基本概念，每輪用 vsetvl 決定 vl，再做載入、計算與累加。

差別在於 attention kernel 的資料路徑更複雜

• Q 來自目前 token
• K/V 可能來自 KV cache
• token 位置 需要透過 req_to_token 間接索引
• BF16 / FP16 需要轉成 FP32 累加
• INT8 KV cache需要縮放係數
• softmax 要維持數值穩定

所以後面看 attention kernel 時，重點不會只放在某一條 RVV 指令。資料配置、分塊大小、cache 重用、執行緒層級平行化和 reduction 都得一起看。

和 SGLang 整合時要處理哪些事

把 RVV intrinsic 寫好只是其中一段，要讓它真的成為 SGLang 後端，還要處理幾個整合點。

第一，後端註冊。--attention-backend rvv 要能建立 RVVAttnBackend。

第二，能力檢查。啟動時要檢查 CPU 是否支援 RVV，也要檢查 sgl-kernel 裡是否有對應運算子。

第三，當 kernel 不存在或 cross-attention 尚未支援時，要退回 Torch Native 或給清楚錯誤。INT8 KV scale 還要再分情況：floating-point K/V 可以用動態 per-token scale 量化後繼續走 RVV；已量化的 INT8 K/V 缺少對應 scale 時則應直接報錯。

第四，tensor資料配置。Python 端要把 q、k、v、KV cache buffer 、請求映射 tensor 整理成 C++ kernel 能讀的形狀。

第五，自訂運算子結構定義。C++ 註冊時要標出哪些 tensor 會被寫入，避免 PyTorch 對 alias 和副作用做錯判斷。

第六，效能測試和測試。Issue #18072 的 交付項目提到單元測試、Qwen2.5-1.5B-Instruct GSM8K testing，以及 RVV vs Torch native 效能比較。

今天先走到這裡

今天先把我的 branch 上的 SGLang RVV 後端路線接起來

• SGLang 是 LLM / multimodal 模型推論服務框架，執行環境會根據後端選擇不同 attention implementation。
• RISC-V CPU 需要 CPU 原生 RVV 路徑，才能利用 C++ RVV intrinsic kernel 加速。
• RVVAttnBackend 負責 Python 端分派、fallback、中繼資料 buffer 、decode / extend 參數整理。
• torch_extension_riscv64.cpp 把 C++ 函式 註冊成 torch.ops.sgl_kernel.*。
• decode.cpp 和 extend.cpp 裡會進入實際的 RVV intrinsic 迴圈。

參考資料

• SGLang Documentation: Welcome to SGLang
• Lianmin Zheng et al., “SGLang: Efficient Execution of Structured Language Model Programs,” NeurIPS 2024
• SGLang paper (arXiv:2312.07104)
• SGLang issue #18072: Adding RVV Attention Backend
• SGLang PR #22193: Add RVV kernels in sgl-kernel
• SGLang PR #22219: Adding RVV Attention Backend for SGLang
  參考 Repo：
• python/sglang/srt/layers/attention/rvv_backend.py、
• sgl-kernel/csrc/cpu/riscv64/torch_extension_riscv64.cpp、
• sgl-kernel/csrc/cpu/riscv64/decode.cpp、
• sgl-kernel/csrc/cpu/riscv64/extend.cpp