跳至主要内容

Runtime 架構

TinyLoop 採用刻意收斂的架構。若拆成模型載入、推論編排、生成、kernel 與驗證層，會更容易理解。

1. 模型載入​

主要實作：

  • src/model.cpp
  • include/model.h

責任：

  • 解析 .tinyloop header
  • 驗證維度與相容性約束
  • 上傳 tensor 到 GPU
  • 配置 runtime buffer
  • 回報權重與 buffer 記憶體

目前重要行為：

  • tensor 可分塊由磁碟串流到 GPU，不需先完整複製到 host
  • logit_row 等可選 buffer 採 lazy 配置
  • 設定 TINYLOOP_EXPERIMENTAL_FP16_BODY=1 時，可在載入時建立 FP16 body cache

2. Buffer Pool​

runtime 會重用少量常駐 buffer：

  • main：FP32 常駐 hidden state
  • norm：FP16 layernorm 輸出
  • scratch1：共享 QKV 或 tiled-MLP 工作空間
  • scratch2：可選 FP16-body MLP 暫存
  • attn_down：共享 attention-output / down-projection buffer
  • logit_row：延遲配置的最後列 logits buffer

此 buffer 策略是 TinyLoop 記憶體輪廓的核心。

3. 推論編排​

主要實作：

  • src/inference.cpp

責任：

  • embedding 路徑選擇
  • pre-block 執行
  • 共享 loop-block 執行
  • 最終 norm
  • logits 物化
  • hidden-trace 與 profiling helper

核心執行形狀：

embed
-> pre.0
-> pre.1
-> loop block x L
-> final norm
-> optional logits

重要切分：

  • 完整 scoring 走全序列路徑
  • 生成類工作負載可使用 score_last_token(...)

4. KV Cache​

主要實作：

  • include/kv_cache.h
  • cuda/kv_cache.cu
  • src/inference.cpp
  • src/generate.cpp

目前能力：

  • prefill 階段填充快取
  • 以已快取 K/V 做單 token decode
  • 預設開啟 cached generation
  • 透過 cache_window 做 sliding-window cache
  • prefix-cache 重用
  • 共享快取的推測式 runtime

目前缺少：

  • paged attention
  • KV-cache 量化
  • 更偏 serving 的 batching 行為

5. 生成層​

主要實作：

  • src/generate.cpp

責任：

  • prompt 驗證
  • sampling
  • EOS 處理
  • prefix-cache helper
  • 推測式 draft/verify 編排

目前生成切分：

  • generate(...)：一般自回歸輸出
  • generate_speculative(...)：自我推測式解碼
  • prefix-cache helper：重複共享前綴

6. CUDA Kernels​

主要實作：

  • cuda/int2_gemm.cu
  • cuda/ops.cu
  • cuda/attention.cu
  • cuda/kv_cache.cu

目前 kernel 分組：

  • 量化 GEMM
  • FP16 GEMM fallback/快速路徑
  • 融合 SwiGLU 與 residual 運算
  • 完整 prefill attention
  • cached 單查詢 decode attention
  • KV-cache append/materialization helper
  • 旋轉位置編碼（RoPE）— cuda/rope.cu

目前 attention 狀態：

  • 直接 cached 單查詢 decode attention 已就位
  • 較安全的 tiled prefill attention 路徑已存在
  • 完整 FlashAttention-2 整合仍是未完成 roadmap 項目
  • GQA head 映射（KV heads 少於 Q heads 的模型）
  • 透過 ring-buffer KV cache 的滑動視窗（cache_window > 0）
  • 虛擬記憶體頁面的 paged attention（TINYLOOP_PAGED_KV=1）

7. 驗證層​

主要實作：

  • tests/
  • CMakeLists.txt

此層很重要，因為 TinyLoop 仍快速演進，效能工作很容易被過度主張。

目前驗證包含：

  • CUDA unit 測試
  • decode parity 測試
  • generation regression
  • tokenizer-aware regression
  • eval-slice regression
  • hot-op microbenchmark

架構原則​

TinyLoop 要保持價值，必須維持其主張：

  • 單一模型家族
  • 明確 runtime 取捨
  • 狹窄公開介面
  • 可量測行為

若太早嘗試成為通用 runtime，就會失去其核心專精。