1 輸入prompt 得到input id

input id： [B,L]?? # batch size , lenth

2 embeding之后得到 input_embeds: [B,L,D]?? # demensions

3 進入Transformer層

先通過linear層得到shape不變的 QKV

多頭注意力 分割Dimension ， kv變成 [B,H,L,head_dim] h是head，head_dim * head = Dimension

(如果有kvcache的話)：

new key [B, H, 1 head_dim]

past key [B,H,L_prev,dim]

concat:[B,H,L_prev+1,dim]

Attn輸出 [B,L,D]

outputs = self.transformer(input_ids, ...)

final_hidden_states = outputs[0]? # shape: [B, L, D]

logits = self.lm_head(final_hidden_states)? # shape: [B, L, vocab_size]

self.transformer 是由多層 Qwen2Block（帶 attention + feedforward）堆疊組成，每層更新一次 hidden state。最終最后一層輸出的 hidden state 就是 final_hidden_states。

logits = lm_head(final_hidden_states)

? → [B, L, vocab_size]

（值是每個位置對每個詞的預測得分）

遇到到是哪個token_id之后 用tokenizer.decode()得到最終的詞

lm_head是什么

語言建模頭（Language Modeling Head)

一個線性層，將輸出的hidden_states映射到詞表維度

# vocab_size = 151936, hidden_size = 4096

self.lm_head = nn.Linear(4096, 151936, bias=False)

得到logits之后還有一個sampling的過程

greedy sampling 直接取最大值

next_token = ops.argmax(next_token_logits, axis=-1)? # [B]

Top-k / Top-p Sampling：

# 用 softmax 得到概率

probs = ops.softmax(next_token_logits, axis=-1)? # [B, vocab_size]

# 根據策略采樣

next_token = sample_from(probs, top_k=50, top_p=0.95, temperature=1.0)

top_k 取前k個概率大的

top_p 累積概率小于0.95