simple-llm-example

Example output

=== 输入序列 ===
ID=(vocab) [5, 0, 9, 5, 1]
tokens= ['小', '狗', '和', '小', '貓']

===== Prefill: 整段序列一次性前向 =====
--- X( All input tokens afeter embedding) | shape=(5, 4) ---
[[0.522 0.297 0.250 0.346]
 [0.100 -0.503 -1.551 0.069]
 [0.963 0.413 0.822 1.897]
 [0.522 0.297 0.250 0.346]
 [-1.062 0.474 -0.919 1.550]]
--- Q | shape=(5, 4) ---
[[0.156 -0.669 0.093 0.711]
 [0.912 -0.868 -0.130 0.449]
 [0.455 -3.413 -2.377 0.335]
 [0.156 -0.669 0.093 0.711]
 [0.168 -3.428 -2.188 -1.698]]
--- K | shape=(5, 4) ---
[[-0.438 0.024 -0.951 -0.424]
 [-0.036 -0.047 1.619 0.121]
 [-1.958 -0.253 -2.934 1.874]
 [-0.438 0.024 -0.951 -0.424]
 [1.340 -0.995 1.122 3.351]]
--- V | shape=(5, 4) ---
[[-0.272 -1.341 -0.154 -0.288]
 [-1.353 0.566 0.898 0.319]
 [-2.124 -3.052 -0.725 0.663]
 [-0.272 -1.341 -0.154 -0.288]
 [-2.857 -1.078 -2.110 3.305]]
--- Scores (masked) | shape=(5, 5) ---
[[-0.237 -INF -INF -INF -INF]
 [-0.243 -0.074 -INF -INF -INF]
 [0.919 -1.832 3.787 -INF -INF]
 [-0.237 0.132 0.461 -0.237 -INF]
 [1.323 -1.796 1.888 1.323 -2.254]]
--- Causal mask | shape=(5, 5) ---
[[1.000 0.000 0.000 0.000 0.000]
 [1.000 1.000 0.000 0.000 0.000]
 [1.000 1.000 1.000 0.000 0.000]
 [1.000 1.000 1.000 1.000 0.000]
 [1.000 1.000 1.000 1.000 1.000]]
--- Attention weights | shape=(5, 5) ---
[[1.000 0.000 0.000 0.000 0.000]
 [0.458 0.542 0.000 0.000 0.000]
 [0.054 0.003 0.943 0.000 0.000]
 [0.183 0.265 0.368 0.183 0.000]
 [0.261 0.012 0.459 0.261 0.007]]
--- Attn out | shape=(5, 4) ---
[[1.081 -0.239 -1.203 -0.949]
 [-0.457 1.189 -0.502 0.402]
 [1.586 1.862 -4.599 -3.369]
 [0.531 1.276 -2.189 -1.237]
 [1.290 0.844 -2.881 -2.139]]
--- FFN hidden (ReLU) | shape=(5, 8) ---
[[1.829 1.615 0.000 1.566 1.778 0.000 0.000 2.919]
 [0.000 0.097 0.000 1.725 1.259 0.901 0.000 0.000]
 [3.600 4.912 0.000 4.959 5.207 0.053 0.000 0.314]
 [1.697 2.740 0.000 2.491 2.710 0.115 0.000 0.000]
 [1.285 0.711 0.000 2.711 3.684 1.559 0.000 0.000]]
--- y2 = y1 + FFN | shape=(5, 4) ---
[[2.081 10.511 -7.734 -3.486]
 [-0.006 2.565 -4.608 1.404]
 [-1.564 14.231 -16.209 -4.957]
 [-1.522 7.220 -8.038 -2.649]
 [0.468 8.404 -8.166 0.750]]
最后位置的下一个 token 预测 → '他'

===== Generation: 生成最多 2 个 token =====
========================生成第 1 个 token========================
--- Gen/x (embedding) | shape=(1, 4) ---
[[0.013 1.454 -0.265 2.720]]
--- Gen/Q | shape=(1, 4) ---
[[0.449 -5.690 -2.266 -0.271]]
--- Gen/K_cur | shape=(1, 4) ---
[[0.627 -1.147 -1.246 2.850]]
--- Gen/V_cur | shape=(1, 4) ---
[[-3.914 -4.866 -3.142 3.226]]
--- Gen/K_all (cache) | shape=(6, 4) ---
[[-0.438 0.024 -0.951 -0.424]
 [-0.036 -0.047 1.619 0.121]
 [-1.958 -0.253 -2.934 1.874]
 [-0.438 0.024 -0.951 -0.424]
 [1.340 -0.995 1.122 3.351]
 [0.627 -1.147 -1.246 2.850]]
--- Gen/V_all (cache) | shape=(6, 4) ---
[[-0.272 -1.341 -0.154 -0.288]
 [-1.353 0.566 0.898 0.319]
 [-2.124 -3.052 -0.725 0.663]
 [-0.272 -1.341 -0.154 -0.288]
 [-2.857 -1.078 -2.110 3.305]
 [-3.914 -4.866 -3.142 3.226]]
--- Gen/Scores | shape=(1, 6) ---
[[0.969 -1.725 3.350 0.969 1.406 4.429]]
--- Gen/Attention weights | shape=(1, 6) ---
[[0.022 0.001 0.234 0.022 0.033 0.688]]
========================生成第 2 个 token========================
--- Gen/x (embedding) | shape=(1, 4) ---
[[0.013 1.454 -0.265 2.720]]
--- Gen/Q | shape=(1, 4) ---
[[0.449 -5.690 -2.266 -0.271]]
--- Gen/K_cur | shape=(1, 4) ---
[[0.627 -1.147 -1.246 2.850]]
--- Gen/V_cur | shape=(1, 4) ---
[[-3.914 -4.866 -3.142 3.226]]
--- Gen/K_all (cache) | shape=(7, 4) ---
[[-0.438 0.024 -0.951 -0.424]
 [-0.036 -0.047 1.619 0.121]
 [-1.958 -0.253 -2.934 1.874]
 [-0.438 0.024 -0.951 -0.424]
 [1.340 -0.995 1.122 3.351]
 [0.627 -1.147 -1.246 2.850]
 [0.627 -1.147 -1.246 2.850]]
--- Gen/V_all (cache) | shape=(7, 4) ---
[[-0.272 -1.341 -0.154 -0.288]
 [-1.353 0.566 0.898 0.319]
 [-2.124 -3.052 -0.725 0.663]
 [-0.272 -1.341 -0.154 -0.288]
 [-2.857 -1.078 -2.110 3.305]
 [-3.914 -4.866 -3.142 3.226]
 [-3.914 -4.866 -3.142 3.226]]
--- Gen/Scores | shape=(1, 7) ---
[[0.969 -1.725 3.350 0.969 1.406 4.429 4.429]]
--- Gen/Attention weights | shape=(1, 7) ---
[[0.013 0.001 0.139 0.013 0.020 0.408 0.408]]
生成的 tokens: ['他', '他']