Aphasia

1import numpy as np
2
3# arange(stop, dtype) 生成 [0, stop) 的等差序列；reshape(3, 4) 将 12 个元素组织为 3 行 4 列。
4# reshape 在 strides 兼容时返回视图，否则可能复制或报错。
5a = np.arange(12, dtype=np.int64).reshape(3, 4)
6
7print(a.shape)    # (3, 4)
8print(a.strides)  # (32, 8)

1offset(i, j) = i * strides[0] + j * strides[1]
2             = i * 32 + j * 8

 1ndarray object
 2┌───────────────────────────┐
 3│ Python Object Header      │
 4├───────────────────────────┤
 5│ Array Metadata            │
 6│   dtype   : float32       │
 7│   shape   : (B,H,S,D)     │
 8│   strides : (s0,s1,s2,s3)│
 9│   data ptr / offset       │
10└──────────────┬────────────┘
11               │
12               ▼
13Data Buffer ──► [0.12 | -0.45 | 0.88 | 1.02 | ...]

1a.T          # shape (4, 3), strides 从 (32, 8) 变为 (8, 32)
2a[:, 1]      # shape (3,), strides (32,), 起点偏移 8 字节
3a[::2]       # 第 0 维 stride 翻倍
4a[::-1]      # 第 0 维使用负 stride

1# arange(6) 生成整数 0 到 5；未显式 dtype 时使用平台默认整数类型。
2x = np.arange(6)
3y = x[1:4]
4y[0] = 100
5print(x)  # [  0 100   2   3   4   5]

1print(a.flags["C_CONTIGUOUS"])
2print(a.flags["F_CONTIGUOUS"])
3# shares_memory(a, b) 判断两个数组是否可能引用同一底层内存，而不是比较数值是否相等。
4print(np.shares_memory(a, a.T))

1(B, S, H, D) -> transpose -> (B, H, S, D)

1a: (3, 1)          b: (1, 4)          a + b: (3, 4)
2┌───┐              ┌───┬───┬───┬───┐  ┌───┬───┬───┬───┐
3│ a │              │ b │ b │ b │ b │  │a+b│a+b│a+b│a+b│
4├───┤              └───┴───┴───┴───┘  ├───┼───┼───┼───┤
5│ a │                                  │a+b│a+b│a+b│a+b│
6├───┤                                  ├───┼───┼───┼───┤
7│ a │                                  │a+b│a+b│a+b│a+b│
8└───┘                                  └───┴───┴───┴───┘

 1# default_rng(seed) 创建独立的 Generator；固定 seed=42 便于复现实验。
 2rng = np.random.default_rng(42)
 3# normal(size=...) 从标准正态分布采样；astype(float32) 把默认 float64 结果转换为 float32。
 4x = rng.normal(size=(1000, 64)).astype(np.float32)
 5
 6# mean(axis=0) 沿样本维求每一列均值；keepdims=True 保留为 (1, 64)，便于广播。
 7mean = x.mean(axis=0, keepdims=True)  # (1, 64)
 8# std 的 axis/keepdims 语义与 mean 相同，返回每一列的总体标准差（默认 ddof=0）。
 9std = x.std(axis=0, keepdims=True)    # (1, 64)
10x_norm = (x - mean) / std
11
12# arange(4, dtype=float32) 生成 [0,1,2,3]；[:, None] 增加长度为 1 的轴，得到 (4, 1)。
13row_weight = np.arange(4, dtype=np.float32)[:, None]  # (4, 1)

 1Python: np.matmul / np.linalg.solve / np.linalg.svd
 2                         │
 3                         ▼
 4NumPy C 层：校验 shape、dtype、广播、布局，准备输入输出 Buffer
 5                         │
 6                         ▼
 7BLAS/LAPACK 接口：GEMM、GESV、GESDD 等标准例程
 8                         │
 9                         ▼
10具体实现：OpenBLAS / MKL / Accelerate / BLIS / Netlib ...
11                         │
12                         ▼
13CPU kernel：SIMD、cache blocking、packing 与多线程

 1import numpy as np
 2
 3# arange(stop, dtype) 创建连续向量；这里显式使用 float64 与后续矩阵 dtype 保持一致。
 4x = np.arange(4, dtype=np.float64)
 5# 第二次 arange 使用相同 stop/dtype，得到与 x shape 和类型一致的向量。
 6y = np.arange(4, dtype=np.float64)
 7# dot(a, b) 对两个一维数组计算内积，对应 BLAS Level 1 的 dot 语义。
 8inner = np.dot(x, y)
 9
10# array(object, dtype) 创建 2×2 稠密矩阵；二维矩阵默认使用 C order。
11A = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float64)
12# array 为 B 创建同样的 2×2 float64 稠密矩阵，满足 matmul 的内维匹配要求。
13B = np.array([[2.0, 0.0], [1.0, 2.0]], dtype=np.float64)
14# matmul(A, B) 对二维输入执行矩阵乘，底层通常进入 BLAS Level 3 GEMM。
15C = np.matmul(A, B)
16# A @ x 是矩阵-向量乘，数学语义对应 BLAS Level 2 GEMV。
17z = A @ x[:2]

1# array([1,0], dtype) 创建 float64 右端向量；linalg.solve(a, b) 求解 a @ x = b。
2# a 必须是方阵，b 可以是一组或多组右端项。
3solution = np.linalg.solve(A, np.array([1.0, 0.0], dtype=np.float64))

1# array(..., order="C") 创建行主序矩阵。
2A_c = np.array([[1.0, 2.0], [3.0, 4.0]], order="C")
3# asfortranarray(a) 在必要时复制为列主序；输入已是 F-contiguous 时直接复用。
4A_f = np.asfortranarray(A_c)
5
6print(A_c.flags["C_CONTIGUOUS"])  # True
7print(A_f.flags["F_CONTIGUOUS"])  # True

1# show_config() 打印 NumPy 构建期依赖，可查看链接的 BLAS/LAPACK 名称与编译信息。
2np.show_config()
3# show_runtime() 打印运行时 SIMD 与 BLAS 线程后端信息；具体字段随构建方式而异。
4np.show_runtime()

1import numpy as np
2
3# array(..., dtype=complex64) 从 Python 复数创建数组，每个元素由两个 float32 组成。
4z = np.array([1.0 + 2.0j, 3.0 + 4.0j], dtype=np.complex64)
5
6real_view = z.real               # stride 8
7imag_view = z.imag               # stride 8，起点多 4 字节
8# view(float32) 只重解释同一缓冲区，不做数值转换；每个 complex64 展开为两个 float32。
9float_view = z.view(np.float32)  # [1., 2., 3., 4.]

 1def apply_rotary_emb_complex(
 2    x: np.ndarray,
 3    base: float = 10000.0,
 4) -> np.ndarray:
 5    """x: (seq_len, head_dim), float32，且 head_dim 为偶数。"""
 6    # asarray(x, dtype) 尽量复用输入；仅当类型不匹配时转换为 float32。
 7    x = np.asarray(x, dtype=np.float32)
 8    seq_len, head_dim = x.shape
 9
10    if head_dim % 2 != 0:
11        raise ValueError("head_dim must be even")
12    if x.strides[-1] != x.itemsize:
13        # ascontiguousarray 仅在布局不连续时复制，确保最后一维可按复数对读取。
14        x = np.ascontiguousarray(x)
15
16    # arange(start, stop, step, dtype) 取得偶数维索引 0,2,...，每两个实数维度组成一个复数维度。
17    dim = np.arange(0, head_dim, 2, dtype=np.float32)
18    inv_freq = 1.0 / (base ** (dim / head_dim))
19    # arange(seq_len) 生成 token 位置；显式 float32 避免后续提升为 float64。
20    positions = np.arange(seq_len, dtype=np.float32)
21    # outer(a, b) 计算外积，得到 (seq_len, head_dim/2) 的旋转角矩阵。
22    angles = np.outer(positions, inv_freq)
23    # exp(1j*angles) 计算单位复数旋转因子；astype(complex64) 压缩到两个 float32。
24    factors = np.exp(1j * angles).astype(np.complex64)
25
26    # view(complex64) 将最后一维相邻两个 float32 零拷贝解释成一个复数。
27    x_complex = x.view(np.complex64)
28    rotated_complex = x_complex * factors
29    # view(float32) 再把复数结果零拷贝展开回实数表示，输出 shape 与 x 相同。
30    return rotated_complex.view(np.float32)

1np.array([1, 2, 3], dtype=np.int32)  # array(object, dtype)：复制/转换 Python 序列为 int32 数组。
2np.asarray(existing_array, dtype=np.float32)  # asarray：类型已匹配时尽量复用，否则转换为 float32。
3np.fromiter((i * i for i in range(10)), dtype=np.int64)  # fromiter(iterable, dtype)：从迭代器单遍构造数组。

1np.eye(3)  # eye(N)：创建 3×3 单位矩阵；可用 k 参数选择偏移对角线。
2np.diag(np.array([1, 2, 3]))  # array 创建向量；diag(v) 把向量放到主对角线。
3np.arange(0, 10, 2)  # arange(start, stop, step)：生成 [0,10) 内步长为 2 的整数。
4np.linspace(2.0, 3.0, 5)  # linspace(start, stop, num)：含两端点地生成 5 个等距值。
5np.logspace(2.0, 3.0, 5)  # logspace(start, stop, num)：生成 10**2 到 10**3 的 5 个对数等距值。

1np.zeros((1024, 64), dtype=np.float32)  # zeros(shape, dtype)：分配并以 0 初始化。
2np.ones_like(x)  # ones_like(a)：创建与 x 的 shape、dtype 默认相同且全为 1 的数组。
3np.full((4, 4), fill_value=-np.inf, dtype=np.float32)  # full：用 -inf 填满指定 shape。
4np.empty((4096,), dtype=np.int64)  # empty：只分配、不初始化；使用前必须完整覆盖。

 1raw = bytearray(4 * 1024)
 2# frombuffer(buffer, dtype) 直接把可缓冲对象解释为数组；默认 count=-1 表示读取全部内容。
 3view = np.frombuffer(raw, dtype=np.float32)
 4
 5# memmap(filename, dtype, mode, shape) 建立文件映射数组；mode="r" 表示只读且不把全文件载入内存。
 6mapped = np.memmap(
 7    "vectors.bin",
 8    dtype=np.float32,
 9    mode="r",
10    shape=(1000, 128),
11)

1a.reshape(3, 4)  # reshape(*shape)：改变逻辑形状；布局兼容时返回视图。
2a.transpose(1, 0)  # transpose(*axes)：按给定轴顺序置换维度，通常只修改 strides。
3a[None, ...]
4a.swapaxes(-1, -2)  # swapaxes(axis1, axis2)：交换最后两个轴，负数轴从末尾计数。
5a.view(np.float32)  # view(dtype)：重解释底层字节，不执行数值类型转换。
6a.ravel()  # ravel(order="C")：拉平成一维并尽量返回视图。

1np.concatenate([x, y], axis=0)  # concatenate：沿已有第 0 轴拼接，其他轴必须兼容；总会物化输出。
2np.stack([x, y], axis=0)  # stack：新增第 0 轴后堆叠，输入 shape 必须完全一致。
3np.tile(x, (2, 3))  # tile(a, reps)：沿两个维度分别重复 2、3 次，会真实复制数据。
4np.repeat(x, 2, axis=0)  # repeat(a, repeats, axis)：把第 0 轴的每个元素重复两次。
5np.pad(x, 1, mode="constant")  # pad(a, pad_width, mode)：每个轴两端各补 1 个常量值。
6a.flatten()  # flatten(order="C")：拉平成一维且总是返回副本。

1a[1:10:2]       # basic slicing，通常是 view
2a[[0, 2, 5]]    # fancy indexing，通常是 copy
3a[a > 0]        # boolean indexing，通常是 copy
4
5a[a < 0] = 0
6positive = np.where(a > 0, a, 0)  # where(condition, x, y)：逐元素按条件选择，支持广播。
7indices = np.nonzero(a > 0)  # nonzero(condition)：返回每个维度的命中下标数组元组。

1x.sum(axis=0, keepdims=True)  # sum(axis=0)：沿第 0 轴归约；keepdims 保留长度为 1 的轴。
2x.argmax(axis=-1)  # argmax(axis=-1)：返回最后一维最大值的位置，而不是最大值本身。
3np.cumsum(x, axis=0)  # cumsum(a, axis)：沿第 0 轴计算包含当前位置的累计和。
4
5np.einsum("bij,bjk->bik", A, B)  # einsum：按下标规则执行批量矩阵乘，j 为归约维。
6np.einsum("bi,bi->b", x, y)  # 对每个 batch 的 i 维做点积，输出 shape 为 (b,)。
7np.einsum("bhsd,bhtd->bhst", q, k)  # 对 d 维缩并，生成注意力分数 (b,h,s,t)。

 1# array(..., dtype=bool) 创建布尔 mask；普通 bool_ 每个元素占 1 字节，并非位压缩。
 2mask = np.array([True, False, True, True, False, True, False, False])
 3# packbits(a) 按每 8 个布尔值压成一个 uint8；默认沿扁平化数组打包。
 4packed = np.packbits(mask)
 5# unpackbits(a) 将字节展开为 bit；切片去掉最后一个字节中可能存在的补齐位。
 6restored = np.unpackbits(packed)[: mask.size]
 7
 8# array(..., dtype=uint8) 明确用单字节无符号整数保存位模式。
 9a = np.array([0b11010100], dtype=np.uint8)
10# 第二个 array 使用相同 dtype，保证 xor 两侧具有一致的字节宽度和位语义。
11b = np.array([0b10110111], dtype=np.uint8)
12# bitwise_xor 逐位异或；bitwise_count 再统计每个整数中置 1 的 bit 数，即汉明距离。
13distance = np.bitwise_count(np.bitwise_xor(a, b))

1# array 从字符串序列推断字符串 dtype；显式 dtype 可控制固定宽度或使用 StringDType。
2names = np.array(["  Alice ", "BOB", "Carol"])
3# strings.lower 逐元素转小写；strings.strip 再移除每个字符串两端空白。
4clean = np.strings.strip(np.strings.lower(names))
5# strings.startswith(a, prefix) 逐元素判断是否以 "a" 开头，返回布尔数组。
6starts = np.strings.startswith(clean, "a")

1# arange(start, stop, dtype=datetime64[D]) 按天生成左闭右开的日期序列。
2days = np.arange("2026-01-01", "2026-02-01", dtype="datetime64[D]")
3weekly = days[::7]
4delta = days - days[0]
5# astype(datetime64[M]) 把日期转换到月粒度；这不是用固定秒数除法计算月份。
6months = days.astype("datetime64[M]")

1np.logical_and(a > 0, a < 10)  # logical_and(x1, x2)：逐元素逻辑与，两个条件可广播。
2np.isfinite(scores)  # isfinite：标记既不是 NaN、也不是正负无穷的元素。
3np.isnan(scores)  # isnan：逐元素检测 IEEE NaN，返回布尔数组。
4np.allclose(actual, expected, rtol=1e-5, atol=1e-7)  # 判断所有元素是否满足绝对/相对误差阈值。

1# default_rng(seed) 使用推荐的 Generator API；相同 seed 与调用顺序可复现。
2rng = np.random.default_rng(42)
3# sqrt(64) 给出缩放分母；normal(loc, scale, size) 采样 64×64 权重；astype 转成 float32。
4w = rng.normal(0.0, 1.0 / np.sqrt(64), size=(64, 64)).astype(np.float32)
5# choice(a, size, replace=False) 从 [0,a) 无放回抽取 128 个整数。
6idx = rng.choice(50_257, size=128, replace=False)
7
8# spawn(n) 从当前 Generator 派生 n 个统计独立的子随机流，适合并行 worker。
9worker_rngs = rng.spawn(4)

1np.exp(x)  # exp：逐元素计算 e**x。
2np.log1p(x)  # log1p：计算 log(1+x)，x 接近 0 时比直接写 log(1+x) 更稳定。
3np.expm1(x)  # expm1：计算 exp(x)-1，小 x 时减少相消误差。
4np.clip(x, lo, hi)  # clip(a, min, max)：把值截断到闭区间 [lo, hi]。
5np.gradient(x)  # gradient：用有限差分估计梯度；默认样本间距为 1。
6
7np.exp(x, out=x)  # out=x 指定写回原数组，避免额外输出分配；会覆盖旧值。

1# unique(a, return_counts=True) 返回排序后的唯一值以及每个值的出现次数。
2values, counts = np.unique(labels, return_counts=True)
3# isin(element, test_elements) 逐元素判断 labels 是否属于 allowed_classes。
4keep = np.isin(labels, allowed_classes)
5clean = X[keep]

1order = np.argsort(scores)  # argsort：返回完整升序排列对应的原始下标。
2topk = np.argpartition(scores, -k)[-k:]  # argpartition(kth=-k)：只保证最大的 k 项位于末尾，内部未排序。
3bucket = np.searchsorted(edges, values)  # searchsorted：在有序 edges 中二分查找每个 value 的插入位置。
4multi_key = np.lexsort((secondary_key, primary_key))  # lexsort(keys)：稳定多键排序，最后一个 key 是主键。

1# hanning(M) 生成长度 256 的 Hann 窗，端点趋近于 0，用于减轻频谱泄漏。
2w = np.hanning(256)
3# fft.rfft(a) 对实数输入做一维 FFT，只返回非负频率的 N/2+1 个复数结果。
4spectrum = np.fft.rfft(signal[:256] * w)
5
6# sliding_window_view(x, window_shape=5) 用 strides 创建长度 5 的重叠窗口视图。
7windows = np.lib.stride_tricks.sliding_window_view(x, 5)
8# mean(axis=1) 对每个窗口归约，得到移动平均；窗口视图本身不复制，结果会分配。
9moving_avg = windows.mean(axis=1)

 1import numpy as np
 2
 3N = 100
 4h = 1.0 / N
 5# linspace(start, stop, num) 在闭区间 [0,1] 生成 N+1 个等距网格点，包含两个边界。
 6x = np.linspace(0.0, 1.0, N + 1)
 7x_inner = x[1:-1]
 8n = N - 1
 9
10A = (
11    # ones(n) 生成主对角线系数；diag(v, k=0) 把它放到主对角线。
12    np.diag(2.0 * np.ones(n))
13    # diag(v, k=1) 构造上偏移一格的对角线；ones(n-1) 与其长度匹配。
14    + np.diag(-1.0 * np.ones(n - 1), k=1)
15    # k=-1 表示下偏移一格的对角线，从而组成三对角矩阵。
16    + np.diag(-1.0 * np.ones(n - 1), k=-1)
17)
18
19# sin 对 x_inner 逐元素计算正弦；np.pi 提供双精度圆周率常量。
20f = (np.pi ** 2) * np.sin(np.pi * x_inner)
21b = (h ** 2) * f
22
23# linalg.solve(A, b) 求解 A @ u = b；要求 A 为方阵且与 b 维度兼容。
24u_inner = np.linalg.solve(A, b)
25# concatenate(seq, axis=0) 把左边界、内部解和右边界拼成完整解，会分配新数组。
26u_numerical = np.concatenate(([0.0], u_inner, [0.0]))
27# sin 逐元素计算解析解，用于和数值结果核对。
28u_exact = np.sin(np.pi * x)
29
30# abs 先计算逐元素绝对误差；max 再归约得到无穷范数误差。
31max_error = np.max(np.abs(u_numerical - u_exact))
32print(f"N={N}, L_inf error={max_error:.3e}")

 1文本
 2  ↓ tokenizer
 3token ids
 4  ↓ embedding
 5隐藏状态 x: (B, S, D)
 6  ↓ Decoder Layer × N      ← 本节只展开其中一层
 7最终 RMSNorm
 8  ↓ LM Head
 9下一 token 的 logits
10  ↓ sampling / argmax
11新的 token

1Prefill（一次输入整个 prompt）
2S = 5, T = 5     Q @ Kᵀ → (5 × 5)，需要因果 mask
3
4Decode（追加一个新 token）
5S = 1, T = 6     Q @ Kᵀ → (1 × 6)，读取已有 K/V
6下一步
7S = 1, T = 7     Q @ Kᵀ → (1 × 7)，继续复用历史 K/V

  1import numpy as np
  2
  3
  4def rms_norm(x: np.ndarray, weight: np.ndarray, eps: float = 1e-6) -> np.ndarray:
  5    # mean(..., axis=-1, keepdims=True) 计算每个 token 最后一维的均方，并保留归约轴用于广播。
  6    variance = np.mean(x * x, axis=-1, keepdims=True)
  7    # sqrt 逐元素开方；eps 防止方差为 0 时除零。
  8    return x * (1.0 / np.sqrt(variance + eps)) * weight
  9
 10
 11def silu(x: np.ndarray) -> np.ndarray:
 12    # exp(-x) 逐元素计算指数，组成 SiLU：x / (1 + exp(-x))。
 13    return x / (1.0 + np.exp(-x))
 14
 15
 16def precompute_freqs_cis(
 17    head_dim: int,
 18    max_seq_len: int,
 19    base: float = 10000.0,
 20) -> np.ndarray:
 21    # arange(start, stop, step, dtype) 取得偶数维索引；两个实数维组成一个复数旋转维。
 22    dim = np.arange(0, head_dim, 2, dtype=np.float32)
 23    inv_freq = 1.0 / (base ** (dim / head_dim))
 24    # arange(max_seq_len) 生成全部 token 位置，显式 float32 控制内存和运算类型。
 25    positions = np.arange(max_seq_len, dtype=np.float32)
 26    # outer(positions, inv_freq) 计算外积，输出 (max_seq_len, head_dim/2) 角度矩阵。
 27    angles = np.outer(positions, inv_freq)
 28    # exp(1j*angles) 生成单位复数；astype(complex64) 把实部、虚部限制为 float32 精度。
 29    return np.exp(1j * angles).astype(np.complex64)
 30
 31
 32def apply_rotary_emb(x: np.ndarray, freqs_cis: np.ndarray) -> np.ndarray:
 33    """x: (B, H, S, HD), float32，且最后一维连续。"""
 34    # asarray(x, dtype) 在输入已经是 float32 ndarray 时复用，否则执行类型转换。
 35    x = np.asarray(x, dtype=np.float32)
 36    if x.shape[-1] % 2 != 0:
 37        raise ValueError("head_dim must be even")
 38    if x.strides[-1] != x.itemsize:
 39        # ascontiguousarray 只在需要时复制，保证相邻两个 float32 在物理上连续。
 40        x = np.ascontiguousarray(x)
 41
 42    # view(complex64) 将最后一维的 float32 两两零拷贝重解释为复数。
 43    x_complex = x.view(np.complex64)
 44    factors = freqs_cis[None, None, :, :]
 45    # 复数乘法完成旋转；view(float32) 再零拷贝展开为原来的实数 shape。
 46    return (x_complex * factors).view(np.float32)
 47
 48
 49class KVCache:
 50    def __init__(
 51        self,
 52        batch_size: int,
 53        num_heads: int,
 54        max_seq_len: int,
 55        head_dim: int,
 56    ) -> None:
 57        shape = (batch_size, num_heads, max_seq_len, head_dim)
 58        # empty(shape, dtype) 仅分配 K Cache，不初始化；后续 append 会覆盖有效区间。
 59        self.k = np.empty(shape, dtype=np.float32)
 60        # V Cache 与 K Cache 使用相同 shape 和 float32 物理类型。
 61        self.v = np.empty(shape, dtype=np.float32)
 62        self.length = 0
 63
 64    def append(
 65        self,
 66        k: np.ndarray,
 67        v: np.ndarray,
 68    ) -> tuple[np.ndarray, np.ndarray]:
 69        step = k.shape[2]
 70        end = self.length + step
 71        if end > self.k.shape[2]:
 72            raise ValueError("KV cache capacity exceeded")
 73
 74        self.k[:, :, self.length:end, :] = k
 75        self.v[:, :, self.length:end, :] = v
 76        self.length = end
 77        return self.k[:, :, :end, :], self.v[:, :, :end, :]
 78
 79
 80class TransformerDecoderLayer:
 81    def __init__(
 82        self,
 83        d_model: int,
 84        num_heads: int,
 85        d_ff: int,
 86        max_seq_len: int,
 87        rng: np.random.Generator,
 88    ) -> None:
 89        if d_model % num_heads != 0:
 90            raise ValueError("d_model must be divisible by num_heads")
 91
 92        self.d_model = d_model
 93        self.num_heads = num_heads
 94        self.head_dim = d_model // num_heads
 95        self.freqs = precompute_freqs_cis(self.head_dim, max_seq_len)
 96
 97        # sqrt(d_model) 计算 Xavier 风格缩放；float32(...) 把标量显式收窄，避免类型提升。
 98        scale = np.float32(1.0 / np.sqrt(d_model))
 99
100        def weight(*shape: int, s: np.float32 = scale) -> np.ndarray:
101            # normal(loc, scale, size) 按指定 shape 采样权重；astype(float32) 转换默认 float64 输出。
102            return rng.normal(0.0, float(s), size=shape).astype(np.float32)
103
104        self.wq = weight(d_model, d_model)
105        self.wk = weight(d_model, d_model)
106        self.wv = weight(d_model, d_model)
107        self.wo = weight(d_model, d_model)
108
109        self.w_gate = weight(d_model, d_ff)
110        self.w_up = weight(d_model, d_ff)
111        # sqrt(d_ff) 生成 FFN 输出层缩放，float32 保持与权重 dtype 一致。
112        self.w_down = weight(d_ff, d_model, s=np.float32(1.0 / np.sqrt(d_ff)))
113
114        # ones(shape, dtype) 初始化 RMSNorm 权重为 1，使初始缩放保持恒等。
115        self.attn_norm = np.ones(d_model, dtype=np.float32)
116        # FFN 前的 RMSNorm 同样用 ones(d_model, float32) 初始化可学习缩放向量。
117        self.ffn_norm = np.ones(d_model, dtype=np.float32)
118
119    def forward(self, x: np.ndarray, cache: KVCache, start_pos: int) -> np.ndarray:
120        B, S, D = x.shape
121        H, HD = self.num_heads, self.head_dim
122
123        residual = x
124        h = rms_norm(x, self.attn_norm)
125
126        # Q：reshape(B,S,H,HD) 拆分 head；transpose(0,2,1,3) 调整为 (B,H,S,HD)。
127        q = (h @ self.wq).reshape(B, S, H, HD).transpose(0, 2, 1, 3)
128        # K：使用相同的 reshape/transpose 参数，使序列维与 head 维交换位置。
129        k = (h @ self.wk).reshape(B, S, H, HD).transpose(0, 2, 1, 3)
130        # V：同样返回 (B,H,S,HD)，与后续注意力矩阵乘的维度约定一致。
131        v = (h @ self.wv).reshape(B, S, H, HD).transpose(0, 2, 1, 3)
132
133        current_freqs = self.freqs[start_pos : start_pos + S]
134        q = apply_rotary_emb(q, current_freqs)
135        k = apply_rotary_emb(k, current_freqs)
136
137        keys, values = cache.append(k, v)
138        total_s = keys.shape[2]
139
140        # swapaxes(-1,-2) 交换 K 的最后两轴得到 K^T；sqrt(float32(HD)) 计算注意力缩放因子。
141        scores = (q @ keys.swapaxes(-1, -2)) / np.sqrt(np.float32(HD))
142
143        if S > 1:
144            # full(shape, fill_value, dtype) 创建全为 -inf 的二维 mask。
145            mask = np.full((S, total_s), -np.inf, dtype=np.float32)
146            # triu(mask, k) 保留第 k 条及以上对角线，其余置 0，形成因果遮罩。
147            mask = np.triu(mask, k=total_s - S + 1)
148            scores += mask[None, None, :, :]
149
150        # max(axis=-1, keepdims=True) 取每行最大值并保留轴，用于稳定 softmax。
151        scores -= np.max(scores, axis=-1, keepdims=True)
152        # exp(..., out=scores) 原地指数化，复用 scores 缓冲区。
153        np.exp(scores, out=scores)
154        # sum(axis=-1, keepdims=True) 计算每行归一化分母，保留轴以便广播除法。
155        scores /= np.sum(scores, axis=-1, keepdims=True)
156
157        attn = scores @ values
158        # transpose 恢复 (B,S,H,HD)，reshape 再合并 H 与 HD 为模型维 D。
159        attn = attn.transpose(0, 2, 1, 3).reshape(B, S, D)
160        h = residual + attn @ self.wo
161
162        residual = h
163        h = rms_norm(h, self.ffn_norm)
164        ffn = (silu(h @ self.w_gate) * (h @ self.w_up)) @ self.w_down
165        return residual + ffn
166
167
168# default_rng(42) 创建可复现、与全局随机状态隔离的 Generator。
169rng = np.random.default_rng(42)
170
171B, D, H, FF = 1, 64, 4, 128
172MAX_SEQ_LEN = 32
173PREFILL_LEN = 5
174
175layer = TransformerDecoderLayer(D, H, FF, MAX_SEQ_LEN, rng)
176cache = KVCache(B, H, MAX_SEQ_LEN, D // H)
177
178# normal(size) 生成标准正态 prompt；astype(float32) 与模型权重类型保持一致。
179prompt = rng.normal(size=(B, PREFILL_LEN, D)).astype(np.float32)
180prefill_out = layer.forward(prompt, cache, start_pos=0)
181print("prefill:", prefill_out.shape, "cache_len:", cache.length)
182
183# size=(B,1,D) 模拟单 token decode 输入；astype 避免默认 float64 扩大带宽。
184current = rng.normal(size=(B, 1, D)).astype(np.float32)
185for step in range(3):
186    current = layer.forward(current, cache, start_pos=cache.length)
187    print("decode:", step, current.shape, "cache_len:", cache.length)