當前筆記順序Engine-Layers(當前l(fā)aynorm.py和attention.py)Qwen3的DecoderLayer概覽(主要關注稠密架構(gòu)設計與張量并行)layernorm.py作用解析就是正常的RMSNorm實現(xiàn)不過值得注意的是在qwen3的實現(xiàn)中會多一層RMSNorm的使用?import torch from torch import nn class RMSNorm(nn.Module): def __init__( self, hidden_size: int, eps: float 1e-6, ) - None: super().__init__() self.eps eps self.weight nn.Parameter(torch.ones(hidden_size)) torch.compile def rms_forward( self, x: torch.Tensor, ) - torch.Tensor: #這里先轉(zhuǎn)float32再轉(zhuǎn)回來轉(zhuǎn)float32是為了計算正確轉(zhuǎn)回來是為了顯存和速度成本 #因為大模型訓練 / 推理時為了節(jié)省顯存、提升速度幾乎都會用 低精度張量比如 float16 或 bfloat16但低精度有個致命問題取值范圍太小計算易出錯。 orig_dtype x.dtype x x.float() var x.pow(2).mean(dim-1, keepdimTrue) x.mul_(torch.rsqrt(var self.eps)) x x.to(orig_dtype).mul_(self.weight) return x torch.compile def add_rms_forward( self, x: torch.Tensor, residual: torch.Tensor, ) - tuple[torch.Tensor, torch.Tensor]: orig_dtype x.dtype x x.float().add_(residual.float()) residual x.to(orig_dtype) var x.pow(2).mean(dim-1, keepdimTrue) x.mul_(torch.rsqrt(var self.eps)) x x.to(orig_dtype).mul_(self.weight) return x, residual def forward( self, x: torch.Tensor, residual: torch.Tensor | None None, ) - torch.Tensor | tuple[torch.Tensor, torch.Tensor]: if residual is None: return self.rms_forward(x) else: return self.add_rms_forward(x, residual) ?attention.py作用解析實現(xiàn)硬件級別的KV緩存管理和注意力計算兩大核心功能。KV_CACHE管理使用TRITON實現(xiàn)注意力計算使用FLASH ATTENION相關庫實現(xiàn)同時在問題6中申明了一下模型輸入的數(shù)據(jù)的變換維度這很重要因為這里的處理不同于正常transformerimport torch from torch import nn import triton import triton.language as tl from flash_attn import flash_attn_varlen_func, flash_attn_with_kvcache from nanovllm.utils.context import get_context #這個裝飾器介紹看問題1 triton.jit #用 Triton JIT 編譯的 kernel用于將 KV tensor也是顯存中 存儲到 GPU 緩存中。 def store_kvcache_kernel( key_ptr, key_stride, value_ptr, value_stride, k_cache_ptr, v_cache_ptr, slot_mapping_ptr, D: tl.constexpr, #Triton中的編譯期常量標記用于標記核函數(shù)中必須在編譯階段確定值的參數(shù)。其值見下個函數(shù) ): #當前線程塊block的 ID并行處理每個 token idx tl.program_id(0) #加載該 token 對應的緩存槽 slot tl.load(slot_mapping_ptr idx) #該位置無效比如padding if slot -1: return key_offsets idx * key_stride tl.arange(0, D) value_offsets idx * value_stride tl.arange(0, D) #加載K/V數(shù)據(jù)GPU 顯存 - GPU 寄存器 key tl.load(key_ptr key_offsets) value tl.load(value_ptr value_offsets) cache_offsets slot * D tl.arange(0, D) #寫入緩存GPU 顯存 - GPU 寄存器 tl.store(k_cache_ptr cache_offsets, key) tl.store(v_cache_ptr cache_offsets, value) def store_kvcache(key: torch.Tensor, value: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, slot_mapping: torch.Tensor): N, num_heads, head_dim key.shape #D是單個token在某一層多頭合并后的總維度。 D num_heads * head_dim #最后一維步長為1確保最后一維連續(xù) assert key.stride(-1) 1 and value.stride(-1) 1 #確保頭的維度步長正確也是確保連續(xù) assert key.stride(1) head_dim and value.stride(1) head_dim #確保Cache的步長正確 assert k_cache.stride(1) D and v_cache.stride(1) D #numel表示元素總數(shù)確保映射表大小匹配 assert slot_mapping.numel() N #[(N,)]是Triton的網(wǎng)格配置grid表示啟動N個線程塊。 store_kvcache_kernel[(N,)](key, key.stride(0), value, value.stride(0), k_cache, v_cache, slot_mapping, D) class Attention(nn.Module): def __init__( self, num_heads, head_dim, scale, num_kv_heads, ): super().__init__() self.num_heads num_heads self.head_dim head_dim #scale就是注意力計算中Softmax的縮放因子核心作用是避免注意力分數(shù)Q?K^T過大導致 Softmax 飽和梯度消失。 self.scale scale self.num_kv_heads num_kv_heads self.k_cache self.v_cache torch.tensor([]) def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor): context get_context() k_cache, v_cache self.k_cache, self.v_cache if k_cache.numel() and v_cache.numel(): store_kvcache(k, v, k_cache, v_cache, context.slot_mapping) if context.is_prefill: if context.block_tables is not None: # prefix cache k, v k_cache, v_cache o flash_attn_varlen_func(q, k, v, max_seqlen_qcontext.max_seqlen_q, cu_seqlens_qcontext.cu_seqlens_q, max_seqlen_kcontext.max_seqlen_k, cu_seqlens_kcontext.cu_seqlens_k, softmax_scaleself.scale, causalTrue, block_tablecontext.block_tables) else: # decode #unsqueeze是因為decode時用上一回生成的單個新token輸入需要增加token序列維度 o flash_attn_with_kvcache(q.unsqueeze(1), k_cache, v_cache, cache_seqlenscontext.context_lens, block_tablecontext.block_tables, softmax_scaleself.scale, causalTrue) return o一些問題1.請介紹一下triton.jit triton.jit 是 Triton 框架的核心裝飾器用于將Python 編寫的函數(shù)編譯為高性能的 GPU 核函數(shù)Kernel替代手動編寫 CUDA C 核函數(shù)的繁瑣過程。 核心特性 自動優(yōu)化Triton 會自動處理 GPU 線程調(diào)度、內(nèi)存訪問優(yōu)化、寄存器分配等底層細節(jié)無需開發(fā)者關注 CUDA 網(wǎng)格 / 塊Grid/Block的手動配置 跨架構(gòu)兼容編譯后的核函數(shù)可在不同代際的 NVIDIA GPU如 Ampere、Hopper上高效運行無需針對不同架構(gòu)適配 Python 語法友好用 Python 語法編寫 GPU 邏輯降低異構(gòu)編程門檻 動態(tài)生成代碼支持編譯期常量、動態(tài)形狀等特性兼顧靈活性與性能。 在示例代碼中store_kvcache_kernel 被該裝飾器修飾后會被編譯為 GPU 核函數(shù)負責將 K/V 數(shù)據(jù)寫入緩存的核心邏輯。 2.為什么說idx tl.program_id(0)是獲取當前核函數(shù)的線程ID它后面不是用來算第幾個偏移嗎難道線程和存儲位置是對應的 1tl.program_id(0) 的含義 Triton 核函數(shù)的執(zhí)行模型是「Grid-Program」網(wǎng)格 - 程序 tl.program_id(dim) 獲取當前核函數(shù)實例在 dim 維度上的索引可理解為「線程 ID」更準確的是「Grid 維度的索引」 示例中 tl.program_id(0) 是一維 Grid 的索引取值范圍是 0 ~ N-1因為啟動核函數(shù)時指定了 [(N, )]即 Grid 大小為 N。 2線程與存儲位置的對應關系 示例中核函數(shù)的設計邏輯是每個線程Program負責處理一個 Token 的 K/V 數(shù)據(jù)存儲N 個 Token 對應 N 個線程。 idx 是第 idx 個線程對應處理第 idx 個 Token 的 K/V 數(shù)據(jù) 線程通過 idx 計算該 Token 的 K/V 數(shù)據(jù)在原始張量中的偏移key_offsets idx * key_stride tl.arange(0,D)再計算該 Token 要存入 Cache 的位置cache_offsets。 簡言之線程 IDidx與 Token 索引一一對應而 Token 索引又對應其存儲位置的偏移因此線程和存儲位置是強綁定的每個線程只處理一個 Token 的存儲。 3.解析cache_offsets slot * D tl.arange(0, D) 這是在計算要把數(shù)據(jù)寫進物理顯存KV Cache的具體內(nèi)存地址。 slot物理槽位號可以理解為大樓里的“房間號”。 D一個 Token 的所有 Attention 頭加起來的總數(shù)據(jù)量即 num_heads * head_dim可以理解為“房間的面積”。 slot * D這就走到了要找的那個房間的門口基礎偏移量。 tl.arange(0, D)推開門給房間里 $0$ 到 $D-1$ 的每一個地磚內(nèi)存單元都打上編號。 合在一起就是準確算出這 $D$ 個數(shù)據(jù)要存進這棟大樓的具體哪些絕對地址中。 4.store_kvcache_kernel[(N, )]的[(N, )]是什么意思 這是 Triton 核函數(shù)的啟動配置表示 以「一維 Grid」啟動核函數(shù)Grid 的大小為 N即啟動 N 個并行的 Program / 線程 (N, ) 是 tuple 類型對應 Grid 的維度一維若為 (N, M) 則是二維 GridN 行 M 列。 示例中 N 是 Token 數(shù)量key.shape[0]啟動 N 個線程每個線程處理一個 Token 的 K/V 存儲與問題 3 的線程 - 存儲位置對應邏輯一致。 5.為什么decode階段q要unsqueeze unsqueeze(1) 是為了匹配 FlashAttention 對 Decode 階段輸入維度的要求核心是補充「序列長度」維度 1Prefill 與 Decode 階段的 Q 維度差異 Prefill 階段預填充處理完整的輸入序列Q 的 shape 通常是 [N, num_heads, head_dim]N 是總 Token 數(shù)隱含序列長度維度 Decode 階段逐 token 生成每次只處理一個 Token自回歸生成Q 的原始 shape 是 [batch_size, num_heads, head_dim]缺少「序列長度」維度seq_len1。 2flash_attn_with_kvcache 的輸入要求 該函數(shù)針對 Decode 階段設計期望 Q 的 shape 包含 seq_len 維度即使 seq_len1即 [batch_size, seq_len, num_heads, head_dim]或簡化為 [N, 1, D]。 示例中 q.unsqueeze(1) 是在第 1 維插入 seq_len1讓 Q 的 shape 從 [N, num_heads*head_dim] 變?yōu)?[N, 1, num_heads*head_dim]匹配函數(shù)的輸入維度要求確保 KV Cache 能正確對齊計算。 6.區(qū)別于transformer的數(shù)據(jù)維度這里是flash_attention的實現(xiàn)導致的我們需要進行補充不然會導致后面的代碼理解錯誤 模型的輸入從來不是 (batch, seq)而是所有序列的 token 拼接成的一維張量 (N,)模型內(nèi)部自始至終保持 2D 的 (N, hidden_size)哪些 token 屬于哪條序列不放在張量形狀里而是放在 cu_seqlens / slot_mapping / block_tables / context_lens 這些元數(shù)據(jù)里。這就是 FlashAttention varlen 格式 分頁 KV cache 的標準做法vLLM 也是這么干的。 1. 數(shù)據(jù)源頭引擎層的一維拼接 引擎根本不構(gòu)造 3D 張量。prepare_prefill 里nanovllm/engine/model_runner.py:129 input_ids.extend(seq[start:end]) —— 把本次 step 里所有序列要處理的 token 順序拼進一個 list最終 torch.tensor 形狀是 (N,)N ∑ 各序列 num_scheduled_tokens 同時用 cu_seqlens_q/k 累計每條序列的邊界[0, len1, len1len2, ...]這就是 varlen 格式的目錄 positions 也是拼接的但取值是每條序列內(nèi)部的絕對位置這樣旋轉(zhuǎn)位置編碼才正確。 decode 階段model_runner.py:172更簡單每條序列只取 seq.last_tokeninput_ids 形狀 (num_seqs,)——每條序列一個 token仍然是 1D。 2. 模型內(nèi)部為什么一直是 (N, hidden) embed_tokens(input_ids)nanovllm/models/qwen3_moe.py:269對 1D 的 id 做 embedding直接得到 (N, hidden_size)——2D沒有任何時刻出現(xiàn) batch/split 維。此后 48 層里所有算子都是逐 token 的 RMSNorm對最后一維求 mean/var(N, hidden) → (N, hidden) 注意力見下節(jié)用的是 varlen 內(nèi)核 MLP / MoE純逐 token 線性變換。 所以 qwen3_moe.py:149 解包出的兩個值其實是sequence_length N本 step 里跨所有序列的 token 總數(shù)變量名有誤導性見第 6 節(jié)、hidden_dim 2048。MoE 路由本來就是逐 token 獨立決策每個 token 選出自己的 top-8 專家完全不需要序列邊界信息后續(xù) index_add_(0, top_x, ...)qwen3_moe.py:187按 flat 行號把各專家的貢獻寫回對應 token 行——一維扁平布局恰好就是 index_add_ 想要的形態(tài)。 3. 注意力varlen 內(nèi)核 分頁 KV cache q/k/v 在模型側(cè)被 view 成 (N, num_heads, head_dim)qwen3_moe.py:81-83然后 prefillnanovllm/layers/attention.py:64-70flash_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_k, ...)。flat 張量 cu_seqlens內(nèi)核內(nèi)部按序列邊界做 causal mask每條序列只attend自己的歷史 token。命中前綴緩存時block_tables is not Nonemodel_runner.py:162k/v 直接換成緩存里的整段張量cu_seqlens_k cu_seqlens_q 表示query 是新的、key 有更多歷史。 decodeattention.py:71-74flash_attn_with_kvcache(q.unsqueeze(1), k_cache, v_cache, cache_seqlenscontext_lens, block_tableblock_tables)。q 變成 (bs, 1, heads, dim)內(nèi)核按每條的 context_lens 從分頁緩存里取它自己的歷史 k/v 做 attention輸出 (bs, heads, dim)。 KV cache 的形狀是 (2, num_layers, num_blocks, block_size, num_kv_heads, head_dim)model_runner.py:115每個 token 有固定槽位 slot block編號 × block_size 塊內(nèi)偏移。attention.py:10-30 的 triton kernel 把每步的 k/v 寫進這些槽位if slot -1: return 是給 warmup/CUDA-graph 靜態(tài)布局用的哨兵。這樣 decode 時任何 token 都能按 block_table 找回自己之前所有 token 的 KV——這就是分頁。 4. 收尾logits 和采樣 最后 norm 完還是 (N, hidden)過 lm_head 得 (N, vocab)。但 prefill 階段我們只需要每條序列最后一個 token 的 logits 來采下一個 token所以 ParallelLMHeadnanovllm/layers/embed_head.py:56-66用 cu_seqlens_q[1:] - 1 把 (N, hidden) 收窄成 (num_seqs, hidden)decode 時 N num_seqs天然就是 (num_seqs, vocab)。Sampler 出 (num_seqs,) 的 token idscheduler.postprocessnanovllm/engine/scheduler.py:81逐條 append_token進入下一輪 decode。 5. 為什么非要用這種布局三個原因 零 padding 浪費序列長度不一如果按 (batch, max_len) 打包短序列要補 pad token多余的矩陣乘法全部白算扁平拼接只對真實 token 付費。 chunked prefill 的天然載體scheduler.schedulescheduler.py:42-46允許一條超長序列被拆成多次 stepnum_scheduled_tokens 截斷每次 step 的 flat N 都不同batch的形狀每步可變但永遠是 1D——扁平 layout 使每次 step 的輸入構(gòu)造變成無腦拼接這正是上游 #218 重構(gòu)后的設計。 與 flash-attn / 分頁 cache 的存儲格式配套序列結(jié)構(gòu)信息誰是誰的歷史、KV 放哪全在元數(shù)據(jù)里張量本身不用攜帶這些維。這也是為什么你能看到 prepare_prefill 里對 slot_mapping、block_tables、context_lens 的精心構(gòu)造——它們才是真正的batch 結(jié)構(gòu)。本系列文章(待寫完修正)[1]Nano-VLLM全代碼解析筆記(1)-sequence[2]Nano-VLLM全代碼解析筆記(2)-block_manager[3]Nano-VLLM全代碼解析筆記(3)-llm_engine和scheduler[4]Nano-VLLM全代碼解析筆記(4)-model_runner[5]Nano-VLLM全代碼解析筆記(5)-laynorm和attention[6]Nano-VLLM全代碼解析筆記(6)-embed_head和linear[7]Nano-VLLM全代碼解析筆記(7)-rotary_embedding[8]Nano-VLLM全代碼解析筆記(8)-qwen3與qwen3_moe上一篇[4]Nano-VLLM全代碼解析筆記(4)-model_runner下一篇[6]Nano-VLLM全代碼解析筆記(6)-embed_head和linear