元素级图解 · 2026-11-09

别只背 shape。
看一个数怎样走过 Transformer。

每张图都给出轴的含义、实际小矩阵、变换前后的坐标，以及那一步究竟有没有做计算。从一个格子出发，再回到真实模型。

先走这条最短路线

先做「分头与转置」→「attention 两次乘法」→「合头」。每次只追踪一个格子；能说出它的新坐标，再看整张量。页面较长，可以按下面目录分次完成。

0 · 一个维度就是一个需要选择的坐标

X[b,t,h] 的意思是：先选第 b 条请求，再选第 t 个 token，最后选它向量里的第 h 个数。[B,T,H]=[1,3,4] 表示有1条请求、每条3个 token、每个 token 4个数，总共1×3×4=12个数。方括号记录每个轴的长度，不是元素本身。

B / b：请求数 / 请求下标T / t：位置数 / 位置下标Hq / hq：Q头数 / 头下标d / j：每头宽度 / 分量下标

整页默认小例：B=1、T=3、H=4、Hq=2、Hkv=1、d=2、F=6。分头交互单独用B=2，以便你切换batch；真实模型例子另列。B、T、Hq、d表示轴长；b、t、hq、j表示下标。下标从0开始。各节用独立的数值小例解释一个操作，并非在运行同一个完整模型；投影和attention的静态图与交互使用相同数值。

操作哪些东西改变先问自己
reshape / view轴的分组方式，元素总数不变哪几个相邻逻辑轴被拆开或合并？
transpose / permute轴的顺序与索引关系，值不变新坐标怎样对应旧坐标？
Linear / matmul沿一个轴乘加，计算新值哪一轴被求和？输出保留哪些轴？
split / slice / gather取一段或选若干元素按哪个轴取？是连续区间还是索引列表？
softmax / RoPE / norm通常保持shape，但数值改变在哪个轴计算，哪些位置相互作用？

图中每次只展开二维切片，其他轴明确写成“固定 b=0 / h=1”。四维张量不是必须画成四维立方体；选定两根轴后，可以像普通表格一样读。

1 · Embedding：用整数选行，增加特征轴

  1. 输入 IDs 是 [B,T]，每格是词表中的一个整数索引。
  2. 查表 E 的形状是 [Vocab,H]。输入 token_id=2，就取 E 的第2行，一次拿到H个数。
  3. 每个原来的格子变成一个长度H的向量，所以结果是 [B,T,H]；这不是把 ID reshape 成浮点向量。
X[b,t,h] = E[IDs[b,t],h]
小例 [1,3] → 查 [Vocab,4] 的3行 → [1,3,4]

源码：Embedding / LM head。取行的索引关系由embedding定义；投影是下一步。

Embedding 查表：ID 到向量。逐行展示 ids=[2,0,4] 如何从 E[6,4] 读取三行，得到 X[1,3,4]；输入索引数量与输出浮点元素数不同，查表不是 reshape。
逐行展示 ids=[2,0,4] 如何从 E[6,4] 读取三行，得到 X[1,3,4]；输入索引数量与输出浮点元素数不同，查表不是 reshape。打开原图放大 ↗

2 · QKV：先计算新特征，再切成三段

固定 b=0，输入 X 是3×4矩阵。数学权重 Wqkv 是4×8。每个输出格子都来自一行 X 和一列 W 的4次乘加，因此输出是3×8。这里 Q宽4、K宽2、V宽2，三段不等宽。

[T,H] × [H,C] → [T,C]，C=(Hq+2Hkv)·d=8
Y[t,c] = Σₖ X[t,k]W[k,c]
Q=Y[:,0:4]；K=Y[:,4:6]；V=Y[:,6:8]

数学W和PyTorch weight方向相反：数学用 X @ W，W=[in,out]；Linear.weight=[out,in]，执行的是 X @ weight.T。矩阵图采用前者。

真实 Qwen3-0.6B：H=1024，Hq=16，Hkv=8，d=128；投影为1024→4096，再split为2048、1024、1024。Hq·d=2048，不等于H。小例恰好相等，不能推广。

证据：Qwen3Attention.forward、固定的官方config副本。

QKV 投影与不等宽 split。用完整 X[3,4] 和 W[4,8] 计算 Y[3,8]，将 Q 的4列、K的2列、V的2列分别标色，解释收缩轴和 Linear.weight 的转置约定。
用完整 X[3,4] 和 W[4,8] 计算 Y[3,8]，将 Q 的4列、K的2列、V的2列分别标色，解释收缩轴和 Linear.weight 的转置约定。打开原图放大 ↗

启用 JavaScript 后可点击矩阵与切换步骤；静态图、索引等式和练习在离线时也可阅读。

3 · 分头与转置：先拆一根轴，再交换两根轴

  1. [B,T,Hq·d] → [B,T,Hq,d]：把每个token的一长行切成Hq小段，每段d个数。原末维 c=hq×d+j。
  2. [B,T,Hq,d] → [B,Hq,T,d]：把token轴和head轴交换，便于按每个head做矩阵乘法。值不变，新坐标为 q_heads[b,hq,t,j]=q_split[b,t,hq,j]。
  3. 前一步是拆轴；后一步是换轴。直接reshape到最终shape，可能把不同token和不同head混在同一行。
q_flat[b,t,hq·d+j] = q_split[b,t,hq,j] = q_heads[b,hq,t,j]
reshape 与 transpose 的元素轨迹。编号0至11由 [1,3,4] 拆成 [1,3,2,2]，再转置为 [1,2,3,2]；逐 head 画出矩阵，并追踪编号2的三套坐标。
编号0至11由 [1,3,4] 拆成 [1,3,2,2]，再转置为 [1,2,3,2]；逐 head 画出矩阵，并追踪编号2的三套坐标。打开原图放大 ↗

启用 JavaScript 后可点击矩阵与切换步骤；静态图、索引等式和练习在离线时也可阅读。

进一步：shape相同，stride也可能不同

以交互里的连续输入 [2,3,4] 为例，stride=[12,4,1]。拆末维后 [2,3,2,2] 的stride=[12,4,2,1]；transpose后shape=[2,2,3,2]，stride=[12,2,4,1]。stride表示索引在该轴加1，需要跨过多少个元素的存储位置。

所以 transpose 通常只是改视图，不是把数据物理搬走。contiguous() 在需要时物理重排；reshape 可能返回view或copy。不能教成“非连续一律不能view”：split后的Q虽然行间仍跨着K/V的数据，也可能合法地view拆分末维；要看具体stride条件。

主源：PyTorch Tensor Views 与 Tensor.view 的stride条件。

4 · QKᵀ 和 AV：两次求和，消掉的轴不同

固定一个 batch、一个query head。Q的行是query位置，K的行是key位置；先转置K，把每个key变成一列。Q的一行和Kᵀ的一列内积，产出一个“这对位置有多匹配”的分数。

Q [Tq,d] × Kᵀ [d,Tk] → S [Tq,Tk]：消掉 d
A [Tq,Tk] × V [Tk,d] → O [Tq,d]：消掉 Tk

softmax在每行的Tk轴上归一化，不把行压成一个数，因此 [Tq,Tk] 形状保持不变。AV才沿key位置求和，把内容汇总回每个query的d维向量。

带缓存时Tq和Tk通常不同。例如历史p=3、新chunk c=2：Q=[2,d]、K/V=[5,d]，scores和mask都是[2,5]，输出=[2,d]。mask允许 j≤p+i；第一行可见4列，第二行可见5列。

QKᵀ 与 AV：保留轴和收缩轴。完整展示一个 head 的 Q、K转置、分数、因果softmax、V和输出，区分第一次沿d求和和第二次沿Tk求和，给出可复算数值。
完整展示一个 head 的 Q、K转置、分数、因果softmax、V和输出，区分第一次沿d求和和第二次沿Tk求和，给出可复算数值。打开原图放大 ↗

启用 JavaScript 后可点击矩阵与切换步骤；静态图、索引等式和练习在离线时也可阅读。

5 · GQA：不同Q头共享K/V头，不减少Q头数

设g=Hq/Hkv，配对规则是 kv_head=⌊q_head/g⌋。本页Hq=2、Hkv=1，两组Q都读取KV head0；这是GQA里的MQA特例。一般GQA例如Hq=4、Hkv=2时，Q0/Q1→KV0，Q2/Q3→KV1。各组Q不同，算出的attention权重仍可不同。输出保留Hq个head。

逻辑配对：K_for_q[b,hq,t,j] = K[b, floor(hq/g),t,j]
Q：Hq个头 → 输出O：Hq个头；缓存K/V：各Hkv个头

公式为展示乘法会把K/V写成“按每个Q头配对”的形状。优化kernel可以直接读取共享head，不要求真的repeat成Hq份。repeat_interleave是便于验证的显式实现，不能拿重复后的张量大小计算真实KV cache。

本地调用：Attention.forward；官方算子约定：PyTorch SDPA。

GQA 的 head 共享与广播。Q的两个 [3,2] 头指向同一个 K/V 头；说明共享的是 K/V，不是注意力权重，也不要求物理复制共享张量。
Q的两个 [3,2] 头指向同一个 K/V 头；说明共享的是 K/V，不是注意力权重，也不要求物理复制共享张量。打开原图放大 ↗

6 · 合头：同一个token的各head，要放回同一行

  1. 如果O是[B,Hq,T,d]，先transpose成[B,T,Hq,d]。
  2. 再合并最后两根轴，得到[B,T,Hq·d]。每行收集同一token的所有head。
  3. 经过Wo投影，Hq·d→H，回到[B,T,H]，这时才能与残差输入逐元素相加。
concat[b,t,hq·d+j] = O[b,hq,t,j]
Y[b,t,k] = Σ𝚌 concat[b,t,c] Wo[c,k]

reshape不会替你理解“同一个token”。直接把head-first的O reshape成[B,T,Hq·d]，会把一个head里相邻token拼在一起。上面交互可切到“合头”对照。

回到nano源码要换布局：prefill的FlashAttention直接收发packed [N,Hq,d]，这里没有先出现[B,Hq,T,d]，所以输出可以直接flatten head与d。decode的backend还可能保留长度1轴，flatten仍得到[N,Hq·d]。不要给源码凭空补一个transpose。

证据：qwen3.py 的view / flatten / o_proj、prefill与decode的backend调用。

正确合头与直接 reshape 的反例。同一组12个编号分别按 transpose后合头与直接reshape排列，直观看见后者把不同 token 的值混到一行。
同一组12个编号分别按 transpose后合头与直接reshape排列，直观看见后者把不同 token 的值混到一行。打开原图放大 ↗

7 · SwiGLU：扩宽、切片、逐元素相乘、投回

  1. X=[B,T,H] 经过fused gate_up投影得到[B,T,2F]；小例4→12。
  2. 沿最后一轴split成gate和up，各[B,T,F]；小例每段6列。
  3. SiLU(gate)与up在同一坐标逐元素相乘，还是[B,T,F]。这里的⊙不是矩阵乘法，token轴和特征轴都没有被求和。
  4. down投影用[F,H]，把最后一轴6→4，输出[B,T,H]。残差x+output要求形状一致。
z[b,t,f] = SiLU(gate[b,t,f]) × up[b,t,f]
out[b,t,h] = Σ𝒻 z[b,t,f] Wdown[f,h]

真实Qwen3-0.6B：1024→6144→两份3072→逐元素相乘3072→1024。RMSNorm通常保持形状，只沿指定特征轴计算统计量；Q/K norm沿d，decoder norm沿H。

证据：Qwen3MLP、SiluAndMul。

SwiGLU 的 split、门控和 down 投影。放大一个token的12维gate_up，切为两个6维向量，实际计算SiLU与逐元素乘积，再经6×4矩阵投回hidden宽度。
放大一个token的12维gate_up，切为两个6维向量，实际计算SiLU与逐元素乘积，再经6×4矩阵投回hidden宽度。打开原图放大 ↗

8 · RoPE：维度没变，也要看清哪两个分量配对

本地nano的布局是把最后一轴平分成两半，不是直接把相邻两个数配对。用d=4的单独小例，x=[a,b,c,d]拆为x1=[a,b]、x2=[c,d]，配对是(a,c)、(b,d)。

输入坐标配对输出仍占原来的4个坐标
x[0]=a，x[2]=c第0个旋转角θ₀y[0]=a cosθ₀−c sinθ₀；y[2]=a sinθ₀+c cosθ₀
x[1]=b，x[3]=d第1个旋转角θ₁y[1]=b cosθ₁−d sinθ₁；y[3]=b sinθ₁+d cosθ₁

例如令两角都为90°方便手算，[1,2,3,4]→[−3,−4,1,2]。真实各对子频率通常不同，这里不模拟频率。Q和K各按自己的head数旋转，V保持原值；形状不变。

证据：apply_rotary_emb；继续读第19课。

9 · 变长组批：N是有效token总数，不是B×Tmax

A长3、B长2：补齐表示会有[2,3,H]，其中一行是PAD；packed表示只有[5,H]。需要先选择有效token并串接，不能只靠reshape删除PAD。cu_seqlens=[0,3,5]记录两条请求的边界，positions=[0,1,2,0,1]记录各自的位置。

n = cu_seqlens[request] + local_token_index
Qpacked [N,Hq,d]；K/Vpacked [N,Hkv,d]
每段独立attention，不能把5个token当成同一条长度5的请求

这个小例是完整prefill；prefix命中、chunked prefill和decode里，新query长度与可读取key长度会不同。源码的cu_seqlens_q与cu_seqlens_k必须分别读。

证据：prepare_prefill，varlen attention调用。

去掉 padding 后的 packed token 轴。将长度3和2的两条请求从 padded [2,3,4] 收集到 [5,4]，逐行标出来源，并用 cu_seqlens=[0,3,5] 恢复请求边界。
将长度3和2的两条请求从 padded [2,3,4] 收集到 [5,4]，逐行标出来源，并用 cu_seqlens=[0,3,5] 恢复请求边界。打开原图放大 ↗

启用 JavaScript 后可点击矩阵与切换步骤；静态图、索引等式和练习在离线时也可阅读。

10 · KV：增长的是历史位置轴，分页改变寻址

单请求、单层的逻辑K为[Tcache,Hkv,d]，V相同。新chunk的Knew为[c,Hkv,d]，写入后可读长度为Tcache+c；decode时c=1。缓存保存每个历史位置的K/V，不保存一张不断扩大的attention分数矩阵。

历史K [3,1,2] + 新K [2,1,2] → 逻辑可读K [5,1,2]
同时Qnew [2,2,2] → 每个Q头的scores [2,5]

上面的“+”表示沿token轴追加，不是逐元素相加；真实分页缓存通常直接写预分配槽位，不必每次分配新tensor并cat整个历史。

nano的单个K池可理解为[num_blocks,block_size,Hkv,d]，V另有同形状池；还要在外层区分layer与K/V。给定token位置t，先查逻辑块t//block_size，再取物理块号，最后找块内偏移。不同backend可能有不同物理布局，不能把这张nano图推广到所有vLLM或DeepInfer内核。

证据：缓存分配与store_kvcache / slot_mapping。

逻辑时间轴到分页 KV 物理坐标。T=5、block_size=2、block_table=[3,0,4] 时逐token计算逻辑块、偏移、物理块和slot，并将每个K分量填进物理池的格子。
T=5、block_size=2、block_table=[3,0,4] 时逐token计算逻辑块、偏移、物理块和slot，并将每个K分量填进物理池的格子。打开原图放大 ↗

11 · 张量并行：拼接与求和发生在不同位置

本图只演示普通Linear的两路分片，不表示单张4070 Ti正在跑双卡，也不把Hkv=1的小例当成nano的TP=2合法配置。

  1. 按数学W的输出列切分：每卡计算一部分输出特征。若需要完整输出，沿特征轴cat；某些路径把分片继续传给下一层。
  2. 按数学W的输入行切分：每卡计算同形状的部分和。两份partial输出要all-reduce相加，不能cat。
  3. 转成PyTorch weight=[out,in]后，图中的列/行方向会反过来；先确认约定，再看代码里的dim。

nano对Q头与KV头都有整除检查；单卡本课的模型几何不涉及跨卡通信。

继续对照第05课；这张图给你判断cat还是sum的数值依据。

列并行 concat 与行并行 sum。用完整小矩阵展示输入复制、输出列分片、concat，以及下游输入行分片产生同形partial后求和；权重采用数学[in,out]约定。
用完整小矩阵展示输入复制、输出列分片、concat，以及下游输入行分片产生同形partial后求和；权重采用数学[in,out]约定。打开原图放大 ↗

12 · 选最后有效位置，再投影到词表

仍用packed长度[3,2]。最终hidden是[N,H]=[5,4]；最后有效行索引是cu_seqlens[1:]−1=[2,4]。先gather这两行，得到[B,H]=[2,4]，再乘词表投影。

packed行号属于谁用于下一token采样
0 / 1 / 2A的三个位置选第2行hidden，预测A的下一token
3 / 4B的两个位置选第4行hidden，预测B的下一token
hidden [5,4] → gather([2,4]) → [2,4]
[B,H] × Wlm [H,Vocab] → logits [B,Vocab]
按词表轴softmax/采样 → next_ids [B]

gather从5行里选2行；Linear把每行的4个特征投到Vocab个分数；采样最后才把每行分数变成一个token ID。三步改变维度的原因各不相同。

这是nano的prefill优化路径，见ParallelLMHead.forward。通用教学图[B,T,Vocab]表示所有位置的logits，实际生成不必物化每个位置的大词表输出。

闭卷检查：只留下形状，也能说明理由吗？

  1. q_flat[0,1,2]分头并转置后，坐标是什么？
  2. 历史长度7、新chunk长度2、每头d=4：每个head的scores和输出分别是什么形状？
  3. 为什么[2,3,4]不能靠reshape变成[5,4]？
先回答，再展开核对

①Hq=2,d=2时对应q_heads[0,1,1,0]；因为2=1×2+0。②scores=[2,9]、输出=[2,4]；d先被收缩，key位置轴再被收缩。③24个元素不能仅靠reshape变成20个，需先选择有效token去掉PAD。

本地打印每一步

cd ~/vllm-learn
python3 labs/dimension_lab.py

不依赖PyTorch或NumPy，打印索引变换并验证正确/错误反例。这里验证的是维度和数值逻辑，不测GPU性能。

把“本页第几步 + 你选中的那个格子 + 你认为的新坐标”发给我，可以继续拆到一条索引等式。