1
def generate(idx, max_new_tokens):
2
    for _ in range(max_new_tokens):
3
        idx_cond = idx[:, -block_size:] # 逗号是Pytorch分维度写法，表示取所有行，最后block_size列
4
        logits, _ = forward(idx_cond) # [B, T, C]
5
        logits = logits[:, -1, :]
6
        probs = torch.softmax(logits, dim=-1)
7
        next_id = torch.argmax(  # 永远取最大
8
            probs,
9
            dim=-1,
10
            keepdim=True
11
        )
12
        idx = torch.cat((idx, next_id), dim=1) # Pytorch API, dim=1表示按列拼接
13
    return idx # [B, T+max_new_tokens]

1
probs = torch.tensor([0.1, 0.2, 0.7])
2
next_id = torch.multinomial(probs, num_samples=1) # Pytorch API 返回 0、1、2 的概率分别约为 10%、20%、70%

1
w 的 Q
2
  去匹配
3
h、e、l、l、o、w 的 K
4
  然后读取
5
h、e、l、l、o、w 的 V

1
Token 的 QKV、MLP、LayerNorm：T 个 → 1 个
2
Attention scores：T×T → 1×T

1
class MultiHeadAttention:
2
    # ... 省略
3
    def forward(self, x, past_kv=None):
4
        # Prefill/训练时：x 是 [B, T, C]
5
        # 使用 KV Cache Decode 时：x 是 [B, 1, C]
6
        q = x @ self.Wq # [B, 1, C] @ [C, C] = [B, 1, C]
7
        k = x @ self.Wk
8
        v = x @ self.Wv
9

10
        B, T, C = x.shape
11
        q = q.reshape(B, T, n_head, head_size)
12
        k = k.reshape(B, T, n_head, head_size)
13
        v = v.reshape(B, T, n_head, head_size)
14

15
        q = q.transpose(1, 2) # 【B, n_head, T, head_size】
16
        k = k.transpose(1, 2)
17
        v = v.transpose(1, 2)
18

19
        past_length = 0
20
        if past_kv is not None:
21
            past_k, past_v = past_kv
22
            past_length = past_k.shape[2]
23
            k = torch.cat((past_k, k), dim=2) # [B, n_head, past_length + 1, head_size]
24
            v = torch.cat((past_v, v), dim=2)
25

26
        scores = q @ k.transpose(-2, -1)
27
        scores = scores / (head_size ** 0.5)
28

29
        # mask = torch.tril(torch.ones(T, T))
30
        if past_kv is None:
31
            mask = torch.tril(torch.ones(T, T, device=x.device)) # device=x.device确保mask和scores在同一个设备上
32
            scores = scores.masked_fill(mask == 0, float('-inf'))
33
        else:
34
            if T != 1:
35
                raise ValueError("When past_kv is provided, T must be 1.")
36
            # 不需要masking，因为每次只生成一个token，scores的shape是[B, n_head, 1, past_length + 1]，只需要计算当前token和之前所有token的注意力分数
37

38
        scores = scores - scores.max(dim=-1, keepdim=True).values
39

40
        exp_scores = torch.exp(scores)
41

42
        weights = exp_scores / exp_scores.sum(
43
            dim=-1,
44
            keepdim=True
45
        )
46

47
        out = weights @ v
48

49
        out = out.transpose(1, 2)
50
        out = out.reshape(B, T, C)
51

52
        out = out @ self.Wo
53
        present_kv = (k, v)
54
        return out, present_kv
55

56

57
def forward(idx, targets=None, past_kvs=None):
58
    past_length = 0
59
    if past_kvs is not None:
60
        if len(past_kvs) != n_layer:
61
            raise ValueError(f"Expected past_kvs to have length {n_layer}, but got {len(past_kvs)}.")
62
        past_length = past_kvs[0][0].shape[2]  # 每一层的KV Cache长度应该是一样的
63

64
    # Decode阶段T就是1
65
    B, T = idx.shape
66

67
    tok_emb = token_embedding_table[idx]
68
    # 变成只取最后一个位置
69
    positions = torch.arange(
70
        past_length,
71
        past_length + T,
72
        device=idx.device,
73
    )
74
    if past_length + T > block_size:
75
        raise ValueError(
76
            "KV Cache length exceeds block_size."
77
        )
78
    pos_emb = position_embedding_table[positions]
79

80
    x = tok_emb + pos_emb
81

82
    present_kvs = []
83
    for layer_idx, block in enumerate(blocks):
84
        layer_past_kv = None
85
        if past_kvs is not None:
86
            layer_past_kv = past_kvs[layer_idx]
87
        x, layer_present_kv = block.forward(x, past_kv=layer_past_kv)
88
        present_kvs.append(layer_present_kv)
89

90

91
    x = layer_norm(x, ln_f_gamma, ln_f_beta)
92

93
    logits = x @ W_lm + b_lm
94

95
    if targets is None:
96
        loss = None
97
    else:
98
        shifted_logits = logits - logits.max(dim=-1, keepdim=True).values
99
        exp_logits = torch.exp(shifted_logits)
100
        probs = exp_logits / exp_logits.sum(dim=-1, keepdim=True)
101

102
        correct_probs = []
103
        for b in range(B):
104
            for t in range(T):
105
                target_id = targets[b, t]
106
                correct_probs.append(probs[b, t, target_id])
107
        correct_probs = torch.stack(correct_probs)
108

109
        loss = -torch.log(correct_probs).mean()
110

111
    return logits, loss, present_kvs
112

113

114
@torch.inference_mode()
115
def generate_with_kv_cache(idx, max_new_tokens, temperature=1.0, top_k=None):
116
    # 由于位置编码是硬编码的，所以 KV Cache 的长度不能超过 block_size，否则位置编码会出错
117
    # 又由于我没写 KVCache 的 Eviction 策略，比如滑动窗口，所以 KVCache 是无脑增长的
118
    # 所以这里限制长度，否则 KV Cache 会随着生成超过 block_size
119
    if idx.shape[1] + max_new_tokens > block_size:
120
        raise ValueError("The total length of idx and max_new_tokens exceeds block_size.")
121

122
````if max_new_tokens < 0:
123
        raise ValueError("max_new_tokens must be non-negative.")
124

125
    if max_new_tokens == 0:
126
        return idx
127

128
    # Prefill阶段
129
    logits, _, past_kvs = forward(idx) # [B, T, C]
130
    next_id = sample_next_token(
131
        logits[:, -1, :],
132
        temperature=temperature,
133
        top_k=top_k
134
    ) # [B, 1]
135
    idx = torch.cat((idx, next_id), dim=1) # [B, T+1]
136

137
    # Decode阶段，每次只有一个Token向量进入神经网络
138
    for _ in range(max_new_tokens - 1):
139
        logits, _, past_kvs = forward(next_id, past_kvs=past_kvs) # [B, 1, C]
140
        next_id = sample_next_token(
141
            logits[:, -1, :],
142
            temperature=temperature,
143
            top_k=top_k
144
        )
145
        idx = torch.cat((idx, next_id), dim=1) # [B, T+2], [B, T+3], ...
146
    return idx # [B, T+max_new_tokens]