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
