HaoyangPing

更新说明

Thu, 16 Jul 2026 00:00:00 GMT

本文记录更新说明，包括本网站维护情况，以及文章发布与修改更新情况等。

2026-09-07

发布 7 篇 llm-algo-leetcode 推理优化笔记：推理结构地基、从 Logits 到 Token：解码采样策略、Speculative Decoding、KV Cache 与推理服务内存管理、模型量化策略、调度策略、推理优化 Benchmark（Task 0/3/4/5/6，部分为占位待完善）。
「llm-algo-leetcode 推理优化笔记」目录更新为 Task 0～6 全量表。

2026-09-05

发布笔记「FlashAttention 1～4 全解汇总」（llm-algo-leetcode 推理优化 Task 2 学习笔记）。
「llm-algo-leetcode 推理优化笔记」目录新增 Task 2 FlashAttention 一行。

2026-08-27

新开「杂集」模块，内设「统计」「英语角」「理财」三个子模块；其中「统计」用于展示 DeepSeek API 使用、运动步数与屏幕时长等日常数据。

2026-08-26

栏目改名：「文章」栏目改为「笔记」，「AIGC」栏目改为「文章」，链接不变。
文章（原 AIGC）板块移除正文开头的 AI 生成内容声明横幅，后续新文章也不再添加。
文章（原 AIGC）板块的「注意力与KV缓存」分类拆分为「FlashAttention」「注意力优化」「KV Cache」三个分类。
修复文章（原 AIGC）板块通知邮件标题被截断的问题。

2026-08-22

修复音乐播放器部分歌曲无法播放的问题。原因为第三方音乐接口服务不稳定，已切换为更稳定的接口服务并修复。

2026-08-21

发布文章 LeetCode | 最长回文子串、LeetCode | 两数相加、LeetCode | 合并两个有序链表。
更新 LeetCode | 双指针内容，新增最长回文子串讲解与代码。

2026-08-20

LeetGPU | 01_Vector_Addition 与 LeetGPU | 02_Matrix_Multiplication 内容更新，改为中文讲解（英文题面与专有名词保留）。
GPU 架构与内存、KV Cache 与优化技术、llm-algo-leetcode 推理优化笔记统一改为中文版本，并更新封面图。

2026-08-19

新增 llm-algo-leetcode Inference Optimization Notes 笔记导航页，汇总本系列各 Task 笔记。
发布文章 KV Cache and Optimization Techniques（Datawhale llm-algo-leetcode 推理优化课程 Task 1-2 笔记）、GPU Architecture and Memory（Task 1-1 笔记）。
修复站内导航及”应募共事”输入密码后音乐中断的问题。
修复页脚”三年之约”图标不显示的问题。
文章表格统一改为与页面同宽，长内容自动换行。
“应募共事”栏目文案调整，仅向招聘相关单位人员开放，并支持点击复制微信号。

2026-08-12

发布文章 LeetCode | 有效的括号、LeetCode | 反转链表。
AIGC 板块文章评论区已启用。

2026-08-10

新开 AIGC 板块。该板块后续将由基于 DeepSeek 模型的 Agent 自动发布文章，相关文章更新不再在本更新说明中逐条阐述。

2026-08-08

发布文章 LeetCode | 寻找两个正序数组的中位数、LeetCode | 搜索旋转排序数组、LeetCode | 在排序数组中查找元素的第一个和最后一个位置、LeetCode | 寻找旋转排序数组中的最小值。

2026-08-04

发布文章 LeetCode | 搜索插入位置、LeetCode | 搜索二维矩阵。
修复文章列表封面预览裁切问题，并恢复封面悬停反馈。
调整文章列表为默认三列网格，每页显示 9 篇。

2026-08-03

发布文章 LeetCode | 双指针、LeetCode | 哈希表、LeetCode | 二叉树的最大深度、LeetCode | 接雨水、LeetCode | 三数之和。
修复文章列表封面裁切问题，封面图现在完整显示。

2026-08-02

发布文章 LeetCode | 最长连续序列、LeetCode | 盛最多水的容器。

2026-08-01

发布文章 LeetCode | 字母异位词分组、LeetCode | list、dict、set 有什么区别？。

2026-07-31

发布文章 LeetCode | 矩阵置零、LeetCode | 无重复字符的最长子串、LeetCode | 二叉树的中序遍历、LeetCode | 相交链表、LeetCode | 移动零、LeetCode | 两数之和。

2026-07-25

发布文章 LeetGPU | 02_Matrix_Multiplication。
LeetGPU | 01_Vector_Addition 更新测试结果。

2026-07-20

发布文章 LeetGPU | 01_Vector_Addition，记录 Vector Addition 的 CUDA、Triton 与 PyTorch 三种实现方式。

2026-07-17

“应募共事”增加本科相关材料
新增”我的”板块，内设”动态”和”相册”两个子模块
“东吴北上赋”文章移至”我的-动态”板块
增加留言功能
导航栏调整：“友链”并入”社交”，新增”留言”入口

llm-algo-leetcode 推理优化笔记

Wed, 19 Aug 2026 02:00:00 GMT

Datawhale

大家好，我是芯缘，是 Datawhale 社区发起的 2026 年 8 月“llm-algo-leetcode 推理优化方向”组队学习活动的运营助教。本页汇总了我在本次活动中的笔记。

目录#

任务	学习内容	学习笔记
Task 0：推理结构地基	Task 0: Attention MHA GQA	- 推理结构地基
Task 1：GPU 硬件与 KV Cache	Task 1-1: GPU Architecture and Memory Task 1-2: KV Cache and Memory Growth	- GPU 架构与内存 - KV Cache 与优化技术
Task 2：Attention 访存瓶颈与 FlashAttention	Task 2-1: FlashAttention Sim Task 2-2: FlashAttention Memory Model	- FlashAttention 1～4 全解汇总
Task 3：解码算法	Task 3-1: Decoding Strategies Task 3-2: Speculative Decoding Task 3-3: Multi-Token Decoding Task 3-4: Decode Scheduling	- 从 Logits 到 Token：解码采样策略 - Speculative Decoding - 调度策略
Task 4：KV Cache 与推理服务内存管理	Task 4-1: KV Cache and Memory Growth Task 4-2: vLLM PagedAttention Task 4-3: SGLang RadixAttention Task 4-4: Prefix Caching and Chunked Prefill Task 4-5: KV Cache Scheduling	- KV Cache 与推理服务内存管理 - 调度策略
Task 5：量化推理与部署	Task 5-1: Quantization Theory and INT4 INT8 Task 5-2: Quantization W8A16 Task 5-3: GPTQ and AWQ Weight Quantization Task 5-4: FP8 and KV Cache Quantization	- 模型量化策略
Task 6：综合项目	Task 6-1: Inference Performance Comparison Task 6-2: Quantized Inference and Deployment Task 6-3: Speculative Decoding Benchmark Task 6-4: Prefix Caching Benchmark	- 推理优化 Benchmark

2026/8/30 分享会内容合集#

LeetCode | 题解汇总

Fri, 31 Jul 2026 17:40:00 GMT

这里整理本站已完成的 LeetCode 题解。

题解汇总#

类别	编号	题目名称
哈希	001	两数之和
哈希	049	字母异位词分组
哈希	128	最长连续序列
双指针	005	最长回文子串
双指针	011	盛最多水的容器
双指针	015	三数之和
双指针	042	接雨水
双指针	283	移动零
栈	020	有效的括号
二分查找	004	寻找两个正序数组的中位数
二分查找	033	搜索旋转排序数组
二分查找	034	在排序数组中查找元素的第一个和最后一个位置
二分查找	035	搜索插入位置
二分查找	074	搜索二维矩阵
二分查找	153	寻找旋转排序数组中的最小值
滑动窗口	003	无重复字符的最长子串
矩阵	073	矩阵置零
链表	002	两数相加
链表	021	合并两个有序链表
链表	160	相交链表
链表	206	反转链表
二叉树	094	二叉树的中序遍历
二叉树	104	二叉树的最大深度

特别篇#

参考资料#

LeetCode 热题 100

推理结构地基

Mon, 07 Sep 2026 05:00:00 GMT

Datawhale

大家好，我是芯缘，是 Datawhale 社区发起的 2026 年 8 月“llm-algo-leetcode 推理优化方向”组队学习活动的运营助教。本文记录了我学习 Task 0：推理结构地基的笔记。

【敬请期待】

参考资料#

Datawhale: Attention MHA GQA

从 Logits 到 Token：解码采样策略

Mon, 07 Sep 2026 04:00:00 GMT

Datawhale

大家好，我是芯缘，是 Datawhale 社区发起的 2026 年 8 月“llm-algo-leetcode 推理优化方向”组队学习活动的运营助教。本文记录了我学习Task 3：解码算法 Decoding Strategies的笔记。

基本原理#

大模型最后一步输出的不是文字，而是对词表中每个 token 的未归一化打分，称为 logits。它的形状通常是 [vocab_size]，比如 32000 个 token 就对应 32000 个分数。

解码采样的目标，就是把这组 logits 转成下一个 token。常见流程是：

logits→截断或平滑→softmax→采样 / 选择 token\text{logits} \rightarrow \text{截断或平滑} \rightarrow \text{softmax} \rightarrow \text{采样 / 选择 token}logits→截断或平滑→softmax→采样 / 选择 token

常见策略#

策略	核心作用	直观影响
Temperature	调整 logits 的分布尖锐程度	控制生成结果更确定(T<1T<1T<1)还是更随机(T>1T>1T>1)
Top-K	只保留概率最高的 KKK 个 token	删除长尾低概率 token
Top-P / Nucleus	保留累计概率达到 ppp 的最小 token 集合	根据分布形状动态决定候选数量

Temperature#

Temperature 在 softmax 之前对 logits 做缩放：

pi=exp⁡(zi/T)∑jexp⁡(zj/T)p_i=\frac{\exp(z_i/T)}{\sum_j \exp(z_j/T)}pi=∑jexp(zj/T)exp(zi/T)

其中 ziz_izi 是第 iii 个 token 的 logit，TTT 是温度系数。

T<1T<1T<1：分布更尖锐，高分 token 更容易被选中，结果更稳定。
T>1T>1T>1：分布更平滑，低分 token 也有更多机会，结果更随机。

Top-K#

Top-K 采样只保留 logits 中分数最高的 KKK 个 token，其余 token 的概率置为 0：

zi′={zi,i∈TopK⁡(z,K)−∞,otherwisez_i'= \begin{cases} z_i, & i\in \operatorname{TopK}(z, K) \\ -\infty, & \text{otherwise} \end{cases}zi′={zi,−∞,i∈TopK(z,K)otherwise

然后再对保留下来的 token 做 softmax 和采样。

Top-P#

Top-P 也叫 Nucleus Sampling。它先按概率从大到小排序，然后保留累计概率刚好超过阈值 ppp 的最小 token 集合：

S=min⁡{S′:∑i∈S′pi≥p}S=\min \left\{S': \sum_{i\in S'} p_i \ge p \right\}S=min{S′:i∈S′∑pi≥p}

和固定候选数量的 Top-K 不同，Top-P 会根据当前概率分布自动调整候选 token 数量：

分布越集中，候选越少；
分布越平缓，候选越多。

Claude Code的文字水印技术#

Claude Code 的文字水印，本质上也是发生在“从 logits 到 token”的采样阶段。

普通采样时，模型先把 logits 变成概率分布，然后用随机数从这个分布里抽一个 token：

logits→softmax→random sample→token\text{logits} \rightarrow \text{softmax} \rightarrow \text{random sample} \rightarrow \text{token}logits→softmax→random sample→token

文字水印的核心想法是： 不往文本里塞隐藏字符，而是把采样时的随机数，换成由 secret key 和上下文共同决定的伪随机数。这样每一步看起来仍然是在随机采样，但长文本累积下来，会出现只有知道 key 才能检测出来的统计模式。

示意代码：

1
def keyed_uniform(secret_key, context_tokens):
2
    """用 secret key + 上下文生成一个稳定的伪随机数。"""
3
    message = f"{secret_key}|{context_tokens}".encode()
4
    digest = hashlib.sha256(message).digest()
5
    value = int.from_bytes(digest[:8], "big")
6
    return value / 2**64
7

8

9
def sample_with_watermark(logits, secret_key, context_tokens):
10
    """水印采样示意：采样仍按概率分布走，但随机数由 key 和上下文决定。"""
11
    probs = softmax(logits)
12
    u = keyed_uniform(secret_key, context_tokens)
13

14
    cumulative = np.cumsum(probs)
15
    return int(np.searchsorted(cumulative, u))

展开收起

只要 secret_key 和 context_tokens 不变，伪随机数 u 就是稳定的。检测时，如果一段文本里 token 的选择反复符合这套带 key 的伪随机规则，就可以判断它很可能来自对应模型。

真实的 SynthID-Text 会更复杂，不只是简单固定一个随机数。它会在候选 token 的采样过程中引入带密钥的打分或排序规则，让水印信号分布在很多 token 选择里。但直觉上可以先抓住这一点：水印不是文本层面的标记，而是采样随机性的标记。

参考资料#

Speculative Decoding

Mon, 07 Sep 2026 03:00:00 GMT

Datawhale

大家好，我是芯缘，是 Datawhale 社区发起的 2026 年 8 月“llm-algo-leetcode 推理优化方向”组队学习活动的运营助教。本文记录了我学习 Task 3：解码算法的笔记。

基本原理#

投机解码（Speculative Decoding）的目标是减少大模型串行解码次数：小模型先生成一段草稿，大模型一次性验证多个 token。

// Step 1: draft model generates γ tokensx^t+1:t+γ∼q(⋅∣x≤t)// Step 2: target model verifies them in one forwardpi=p(x^t+i∣x≤t,x^t+1:t+i−1)qi=q(x^t+i∣x≤t,x^t+1:t+i−1)// Step 3: accept prefix tokens sequentiallyai=min⁡ ⁣(1,piqi)\begin{aligned} &\text{// Step 1: draft model generates } \gamma \text{ tokens} \\ &\hat{x}_{t+1:t+\gamma}\sim q(\cdot\mid x_{\le t}) \\ &\text{// Step 2: target model verifies them in one forward} \\ &p_i=p(\hat{x}_{t+i}\mid x_{\le t},\hat{x}_{t+1:t+i-1}) \\ &q_i=q(\hat{x}_{t+i}\mid x_{\le t},\hat{x}_{t+1:t+i-1}) \\ &\text{// Step 3: accept prefix tokens sequentially} \\ &a_i=\min\!\left(1,\frac{p_i}{q_i}\right) \end{aligned}// Step 1: draft model generates γ tokensx^t+1:t+γ∼q(⋅∣x≤t)// Step 2: target model verifies them in one forwardpi=p(x^t+i∣x≤t,x^t+1:t+i−1)qi=q(x^t+i∣x≤t,x^t+1:t+i−1)// Step 3: accept prefix tokens sequentiallyai=min(1,qipi)

分类#

解码方法可以先按“是否严格一 token 一步”和“是否需要草稿验证”来分：

类别	核心方式	是否一次推进多 token	例子
自回归解码	每次用目标模型生成下一个 token	否	Greedy Search Sampling Beam Search
Multi-Token Decoding	一次解码步骤尝试产出多个 token	是	Speculative Decoding MTP DFlash DSpark
Speculative Decoding	先生成草稿 token 再由目标模型验证	通常是	经典 Draft-Verify MTP DFlash DSpark

比如 Speculative Decoding 中，当 γ=1\gamma=1γ=1 时，主体模型还是自回归解码；当 γ>1\gamma>1γ>1 时，如果一次验证能接受多个草稿 token，整体流程就具有 Multi-Token Decoding 的特点。但这并不要求 draft model 本身也必须是 Multi-Token Decoding：经典 draft model 通常仍是自回归生成草稿，而在 DFlash 中，draft model 本身也采用了并行的 multi-token 生成方式，因此 draft 阶段本身也属于 Multi-Token Decoding。

MTP#

MTP（Multi-Token Prediction）是在模型内部加入未来 token 预测头，让主模型自己生成草稿 token，因此常被看作 self-speculative decoding。

它的优势是不需要额外部署独立 draft model；代价是模型结构和训练目标要配合 MTP 模块。

【待完善】

DFlash#

DFlash 用块扩散模型生成草稿 token，用并行去噪替代传统自回归草稿生成。

它优化的是 draft 阶段：让草稿 token 的生成更并行。

【待完善】

DSpark#

Datawhale

大家好，我是芯缘，是 Datawhale 社区发起的 2026 年 8 月“llm-algo-leetcode 推理优化方向”组队学习活动的运营助教。本文记录了我学习 Task 2： Attention 访存瓶颈与 FlashAttention 的笔记。

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness（2022）#

一种具有 IO 感知能力的快速、显存高效的精确注意力算法。

Algorithm#

Background#

(1) 内存的金字塔结构，HBM 带宽远小于 On-chip SRAM。

(2) SSS 和 PPP 均为 O(N2)O(N^2)O(N2) ，被反复读写 HBM 时，即使 Attention 的计算量为 O(N2d)O(N^2d)O(N2d)，整体容易呈现 Memory-Bound。

Solution#

沿序列长度分块，在 SRAM 中完成局部 score、online softmax 和输出累加，从而避免将完整的 SSS 和 PPP 写入 HBM。

Input: Q,K,V∈RN×d,SRAM size=MOutput: O=softmax⁡(QK⊤)V// Step 1: Set block sizeBc←⌈M4d⌉,Br←min⁡ ⁣(⌈M4d⌉,d)// Step 2: InitializeO←0,ℓ←0,m←−∞// Step 3: Split by sequence lengthQ→{Q1,…,QTr},Qi∈RBr×dK→{K1,…,KTc},Kj∈RBc×dV→{V1,…,VTc},Vj∈RBc×dO→{O1,…,OTr},ℓ→{ℓ1,…,ℓTr},m→{m1,…,mTr}// Step 4: Forward loopfor j←1 to Tc doKj,Vj: HBM → SRAMfor i←1 to Tr doQi,Oi,ℓi,mi: HBM → SRAMSij←QiKj⊤minew←max⁡ ⁣(mi,rowmax⁡(Sij))P~ij←exp⁡ ⁣(Sij−minew)ℓinew←emi−minewℓi+rowsum⁡ ⁣(P~ij)Oinew←emi−minewℓiOi+P~ijVjℓinewOi,ℓi,mi←Oinew,ℓinew,minewOi,ℓi,mi: SRAM → HBMend forend for// Step 5: Returnreturn O\begin{aligned} \textbf{Input: }& Q,K,V\in\mathbb{R}^{N\times d},\quad \text{SRAM size}=M \\[4pt] \textbf{Output: }& O=\operatorname{softmax}(QK^\top)V \\[8pt] &\text{// Step 1: Set block size} \\[2pt] &B_c \leftarrow \left\lceil \frac{M}{4d} \right\rceil,\quad B_r \leftarrow \min\!\left(\left\lceil \frac{M}{4d} \right\rceil,d\right) \\[8pt] &\text{// Step 2: Initialize} \\[2pt] &O \leftarrow 0,\quad \ell \leftarrow 0,\quad m \leftarrow -\infty \\[8pt] &\text{// Step 3: Split by sequence length} \\[2pt] &Q \rightarrow \{Q_1,\dots,Q_{T_r}\},\quad Q_i\in\mathbb{R}^{B_r\times d} \\[4pt] &K \rightarrow \{K_1,\dots,K_{T_c}\},\quad K_j\in\mathbb{R}^{B_c\times d} \\[4pt] &V \rightarrow \{V_1,\dots,V_{T_c}\},\quad V_j\in\mathbb{R}^{B_c\times d} \\[4pt] &O \rightarrow \{O_1,\dots,O_{T_r}\},\quad \ell \rightarrow \{\ell_1,\dots,\ell_{T_r}\},\quad m \rightarrow \{m_1,\dots,m_{T_r}\} \\[8pt] &\text{// Step 4: Forward loop} \\[2pt] &\textbf{for } j \leftarrow 1 \textbf{ to } T_c \textbf{ do} \\[2pt] &\quad K_j,V_j:\text{ HBM }\rightarrow\text{ SRAM} \\[6pt] &\quad \textbf{for } i \leftarrow 1 \textbf{ to } T_r \textbf{ do} \\[2pt] &\quad\quad Q_i,O_i,\ell_i,m_i:\text{ HBM }\rightarrow\text{ SRAM} \\[6pt] &\quad\quad S_{ij}\leftarrow Q_iK_j^\top \\[6pt] &\quad\quad m_i^{\text{new}} \leftarrow \max\!\left(m_i,\operatorname{rowmax}(S_{ij})\right) \\[6pt] &\quad\quad \widetilde{P}_{ij} \leftarrow \exp\!\left(S_{ij}-m_i^{\text{new}}\right) \\[6pt] &\quad\quad \ell_i^{\text{new}} \leftarrow e^{m_i-m_i^{\text{new}}}\ell_i +\operatorname{rowsum}\!\left(\widetilde{P}_{ij}\right) \\[6pt] &\quad\quad O_i^{\text{new}} \leftarrow \frac{ e^{m_i-m_i^{\text{new}}}\ell_iO_i+\widetilde{P}_{ij}V_j }{ \ell_i^{\text{new}} } \\[10pt] &\quad\quad O_i,\ell_i,m_i \leftarrow O_i^{\text{new}},\ell_i^{\text{new}},m_i^{\text{new}} \\[4pt] &\quad\quad O_i,\ell_i,m_i:\text{ SRAM }\rightarrow\text{ HBM} \\[2pt] &\quad \textbf{end for} \\[2pt] &\textbf{end for} \\[6pt] &\text{// Step 5: Return} \\[2pt] &\textbf{return } O \end{aligned}Input: Output: Q,K,V∈RN×d,SRAM size=MO=softmax(QK⊤)V// Step 1: Set block sizeBc←⌈4dM⌉,Br←min(⌈4dM⌉,d)// Step 2: InitializeO←0,ℓ←0,m←−∞// Step 3: Split by sequence lengthQ→{Q1,…,QTr},Qi∈RBr×dK→{K1,…,KTc},Kj∈RBc×dV→{V1,…,VTc},Vj∈RBc×dO→{O1,…,OTr},ℓ→{ℓ1,…,ℓTr},m→{m1,…,mTr}// Step 4: Forward loopfor j←1 to Tc doKj,Vj: HBM → SRAMfor i←1 to Tr doQi,Oi,ℓi,mi: HBM → SRAMSij←QiKj⊤minew←max(mi,rowmax(Sij))Pij←exp(Sij−minew)ℓinew←emi−minewℓi+rowsum(Pij)Oinew←ℓinewemi−minewℓiOi+PijVjOi,ℓi,mi←Oinew,ℓinew,minewOi,ℓi,mi: SRAM → HBMend forend for// Step 5: Returnreturn O

Details#

Block Size设置 关键约束是一次计算时，Kj,VjK_j,V_jKj,Vj、Qi,OiQ_i,O_iQi,Oi 和局部 score block SijS_{ij}Sij 都要放进 On-chip SRAM：

Bcd=O(M),Brd=O(M),BrBc=O(M).B_cd=O(M),\qquad B_rd=O(M),\qquad B_rB_c=O(M).Bcd=O(M),Brd=O(M),BrBc=O(M).

其中 Bcd=O(M)B_cd=O(M)Bcd=O(M) 来自 Kj,VjK_j,V_jKj,Vj 的驻留需求，Brd=O(M)B_rd=O(M)Brd=O(M) 来自 Qi,OiQ_i,O_iQi,Oi 的驻留需求，BrBc=O(M)B_rB_c=O(M)BrBc=O(M) 来自 Sij=QiKj⊤∈RBr×BcS_{ij}=Q_iK_j^\top\in\mathbb{R}^{B_r\times B_c}Sij=QiKj⊤∈RBr×Bc 的临时存储需求。

FA1 让 BcB_cBc 尽量大：

Bc=Θ(Md),B_c=\Theta\left(\frac{M}{d}\right),Bc=Θ(dM),

这样外层 K,VK,VK,V block 数更少，Q,OQ,OQ,O 被重复扫描的轮数也更少。但 BcB_cBc 变大后，SijS_{ij}Sij 的约束会限制 BrB_rBr：

BrBc=O(M)⇒Br=O(d).B_rB_c=O(M) \quad\Rightarrow\quad B_r=O(d).BrBc=O(M)⇒Br=O(d).

因此论文中取：

Bc=⌈M4d⌉,Br=min⁡ ⁣(⌈M4d⌉,d).B_c=\left\lceil\frac{M}{4d}\right\rceil,\qquad B_r=\min\!\left(\left\lceil\frac{M}{4d}\right\rceil,d\right).Bc=⌈4dM⌉,Br=min(⌈4dM⌉,d).

为什么 mmm 初始化为 −∞-\infty−∞？ mmm 记录的是当前已经扫描过的 score 的行最大值。online softmax 在处理第一个 SijS_{ij}Sij block 时，希望新最大值完全由这个 block 决定：

minew=max⁡ ⁣(mi,rowmax⁡(Sij))=rowmax⁡(Sij).m_i^{\text{new}} =\max\!\left(m_i,\operatorname{rowmax}(S_{ij})\right) =\operatorname{rowmax}(S_{ij}).minew=max(mi,rowmax(Sij))=rowmax(Sij).

把 mim_imi 初始化成 −∞-\infty−∞，就相当于告诉算法“当前还没有看过任何 score”。同时旧分母和旧输出的缩放项会自然消失：

emi−minewℓi=e−∞−minew⋅0=0.e^{m_i-m_i^{\text{new}}}\ell_i =e^{-\infty-m_i^{\text{new}}}\cdot 0 =0.emi−minewℓi=e−∞−minew⋅0=0.

这样第一块和后续块可以使用完全相同的更新公式，不需要单独写 first block 的特殊分支。

复杂度分析 计算量仍然由两次矩阵乘主导，总体 FLOPs 仍为：

O(N2d).O(N^2d).O(N2d).

HBM 访问量的变化更关键。所有 K,VK,VK,V block 在外层循环中合计读入一次，代价是 Θ(Nd)\Theta(Nd)Θ(Nd)；每处理一个 K,VK,VK,V block，需要从 HBM 扫描一遍全部 Q,OQ,OQ,O block 以及 m,ℓm,\ellm,ℓ 统计量，主导代价是 Θ(Nd)\Theta(Nd)Θ(Nd)，不是 Θ(Nd+N2)\Theta(Nd+N^2)Θ(Nd+N2)。

这里不会出现每轮 Θ(N2)\Theta(N^2)Θ(N2) 的 HBM 访问，因为局部 SijS_{ij}Sij 和 P~ij\widetilde{P}_{ij}Pij 虽然总体计算量覆盖了 N2N^2N2 个 attention score，但它们只在 SRAM 中短暂生成和消费，并不会作为完整矩阵写入 HBM。因为：

Tc=NBc=Θ(NdM),T_c=\frac{N}{B_c} =\Theta\left(\frac{Nd}{M}\right),Tc=BcN=Θ(MNd),

所以 HBM 访问量为：

Θ(Nd+NdTc)=Θ(NdTc)=Θ(N2d2M).\Theta(Nd+NdT_c) =\Theta(NdT_c) =\Theta\left(\frac{N^2d^2}{M}\right).Θ(Nd+NdTc)=Θ(NdTc)=Θ(MN2d2).

标准 Attention 需要把完整 SSS 和 PPP 写入并读出 HBM，HBM 访问量为：

Θ(Nd+N2).\Theta(Nd+N^2).Θ(Nd+N2).

当片上 SRAM 足够容纳较大的 block，且通常 M≫d2M\gg d^2M≫d2 时，FlashAttention-1 可以显著减少 HBM 访问。注：SRAM容量通常为约 646464 到 164 KiB164\ \text{KiB}164 KiB；FP16/BF16 下约 32K32\text{K}32K 到 82K82\text{K}82K 个元素。

Source Code#

【这一部分内容待完善～】

FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning（2023）#

通过更好的并行性与工作划分实现更快的注意力算法。

Algorithm#

Background#

(1) 矩阵乘由 Tensor Core 加速，而 non-matmul FLOPs 更贵。FA1 的Thread Block 内部的 Non-Matmul 计算偏多。

(2) FA1 的并行划分局限于batch和head，且在短 Batch / 长序列场景下 Occupancy 不高。

Solution#

FA2 保留 FA1 的精确分块注意力思想，

通过内外循环交换把并行维度扩展到序列长度方向，
通过online softmax 更新公式改写实现减少每个 block 内的 non-matmul 操作。

Input: Q,K,V∈RN×d,Br,BcOutput: O=softmax⁡(QK⊤)V,L// Step 1: Split by sequence lengthQ→{Q1,…,QTr},Qi∈RBr×dK→{K1,…,KTc},Kj∈RBc×dV→{V1,…,VTc},Vj∈RBc×dO→{O1,…,OTr},L→{L1,…,LTr}// Step 2: Parallelize over query row blocksparallel for i←1 to Tr doQi: HBM → SRAMO~i(0)←0,ℓi(0)←0,mi(0)←−∞// Step 3: Scan key/value blocksfor j←1 to Tc doKj,Vj: HBM → SRAMSi(j)←QiKj⊤mi(j)←max⁡ ⁣(mi(j−1),rowmax⁡(Si(j)))P~i(j)←exp⁡ ⁣(Si(j)−mi(j))ℓi(j)←emi(j−1)−mi(j)ℓi(j−1)+rowsum⁡ ⁣(P~i(j))O~i(j)←emi(j−1)−mi(j)O~i(j−1)+P~i(j)Vjend for// Step 4: Normalize once and save logsumexpOi←O~i(Tc)ℓi(Tc)Li←mi(Tc)+log⁡ ⁣(ℓi(Tc))Oi,Li: SRAM → HBMend for// Step 5: Returnreturn O,L\begin{aligned} \textbf{Input: }& Q,K,V\in\mathbb{R}^{N\times d},\quad B_r,B_c \\[4pt] \textbf{Output: }& O=\operatorname{softmax}(QK^\top)V,\quad L \\[8pt] &\text{// Step 1: Split by sequence length} \\[2pt] &Q \rightarrow \{Q_1,\dots,Q_{T_r}\},\quad Q_i\in\mathbb{R}^{B_r\times d} \\[4pt] &K \rightarrow \{K_1,\dots,K_{T_c}\},\quad K_j\in\mathbb{R}^{B_c\times d} \\[4pt] &V \rightarrow \{V_1,\dots,V_{T_c}\},\quad V_j\in\mathbb{R}^{B_c\times d} \\[4pt] &O \rightarrow \{O_1,\dots,O_{T_r}\},\quad L \rightarrow \{L_1,\dots,L_{T_r}\} \\[8pt] &\text{// Step 2: Parallelize over query row blocks} \\[2pt] &\textbf{parallel for } i \leftarrow 1 \textbf{ to } T_r \textbf{ do} \\[2pt] &\quad Q_i:\text{ HBM }\rightarrow\text{ SRAM} \\[4pt] &\quad \widetilde{O}_i^{(0)} \leftarrow 0,\quad \ell_i^{(0)} \leftarrow 0,\quad m_i^{(0)} \leftarrow -\infty \\[8pt] &\quad \text{// Step 3: Scan key/value blocks} \\[2pt] &\quad \textbf{for } j \leftarrow 1 \textbf{ to } T_c \textbf{ do} \\[2pt] &\quad\quad K_j,V_j:\text{ HBM }\rightarrow\text{ SRAM} \\[6pt] &\quad\quad S_i^{(j)} \leftarrow Q_iK_j^\top \\[6pt] &\quad\quad m_i^{(j)} \leftarrow \max\!\left(m_i^{(j-1)},\operatorname{rowmax}(S_i^{(j)})\right) \\[6pt] &\quad\quad \widetilde{P}_i^{(j)} \leftarrow \exp\!\left(S_i^{(j)}-m_i^{(j)}\right) \\[6pt] &\quad\quad \ell_i^{(j)} \leftarrow e^{m_i^{(j-1)}-m_i^{(j)}}\ell_i^{(j-1)} +\operatorname{rowsum}\!\left(\widetilde{P}_i^{(j)}\right) \\[6pt] &\quad\quad \widetilde{O}_i^{(j)} \leftarrow e^{m_i^{(j-1)}-m_i^{(j)}}\widetilde{O}_i^{(j-1)} +\widetilde{P}_i^{(j)}V_j \\[2pt] &\quad \textbf{end for} \\[8pt] &\quad \text{// Step 4: Normalize once and save logsumexp} \\[2pt] &\quad O_i \leftarrow \frac{\widetilde{O}_i^{(T_c)}}{\ell_i^{(T_c)}} \\[6pt] &\quad L_i \leftarrow m_i^{(T_c)}+\log\!\left(\ell_i^{(T_c)}\right) \\[6pt] &\quad O_i,L_i:\text{ SRAM }\rightarrow\text{ HBM} \\[2pt] &\textbf{end for} \\[8pt] &\text{// Step 5: Return} \\[2pt] &\textbf{return } O,L \end{aligned}Input: Output: Q,K,V∈RN×d,Br,BcO=softmax(QK⊤)V,L// Step 1: Split by sequence lengthQ→{Q1,…,QTr},Qi∈RBr×dK→{K1,…,KTc},Kj∈RBc×dV→{V1,…,VTc},Vj∈RBc×dO→{O1,…,OTr},L→{L1,…,LTr}// Step 2: Parallelize over query row blocksparallel for i←1 to Tr doQi: HBM → SRAMOi(0)←0,ℓi(0)←0,mi(0)←−∞// Step 3: Scan key/value blocksfor j←1 to Tc doKj,Vj: HBM → SRAMSi(j)←QiKj⊤mi(j)←max(mi(j−1),rowmax(Si(j)))Pi(j)←exp(Si(j)−mi(j))ℓi(j)←emi(j−1)−mi(j)ℓi(j−1)+rowsum(Pi(j))Oi(j)←emi(j−1)−mi(j)Oi(j−1)+Pi(j)Vjend for// Step 4: Normalize once and save logsumexpOi←ℓi(Tc)Oi(Tc)Li←mi(Tc)+log(ℓi(Tc))Oi,Li: SRAM → HBMend for// Step 5: Returnreturn O,L

Details#

通过交换内外层循环带来的收益

(1) 只需要在外循环里内循环外对 FA1 OiO_iOi进行归一化的rescale。

(2) 因为 QQQ 在外循环，所以不需要每次都把临时的O,m,ℓO,m,\ellO,m,ℓ进行FA1那样内外循环读写，而是每次外循环都是赋值0。

(3)在forward过程中维护的用于训练的参数m,ℓm,\ellm,ℓ不需要保存到 HBM，而是只保存 logsumexp：Li=mi+log⁡(ℓi)L_i=m_i+\log(\ell_i)Li=mi+log(ℓi) 注：如果是纯推理的话，m,ℓ,Lm,\ell,Lm,ℓ,L其实都不需要写回HBM，LLL推理中都不需要有，但是m,ℓm,\ellm,ℓ无论训推无论FA version，都需要中途维护。

复杂度分析

总体 FLOPs 仍为：O(N2d).O(N^2d).O(N2d).

如果按 FA2 的 row-block outer loop 粗略估算 HBM 访问量，每个 QiQ_iQi block 会扫描一遍完整 K,VK,VK,V：

Tr=NBr.T_r=\frac{N}{B_r}.Tr=BrN.Tr=NBr,HBM reads of K,V≈Θ(NdTr).T_r=\frac{N}{B_r},\qquad \text{HBM reads of }K,V\approx \Theta(NdT_r).Tr=BrN,HBM reads of K,V≈Θ(NdTr).

HBM 访问估算：

Θ(NdTr)=Θ(N2dBr).\Theta(NdT_r) =\Theta\left(\frac{N^2d}{B_r}\right).Θ(NdTr)=Θ(BrN2d).

Source Code#

包括Work Partitioning Between Warps 【这一部分内容待完善～】

FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision（2024）#

一种利用异步执行与低精度计算实现快速且准确注意力的算法。

FlashAttention 3 的算法架构与 FlashAttention 2 一致，主要针对 Hopper (H100) 架构做了硬件级优化。

各种优化#

FA3没有算法上面的创新，因此暂缓，以后补～ Datawhale教程里面指出：

核心创新 1：WGMMA 异步计算。利用 Warp Group 级指令，让 Tensor Core 在后台异步执行。

核心创新 2：TMA（Tensor Memory Accelerator）。使用硬件级搬运器把数据从全局显存搬到共享内存，释放搬运线程。

核心创新 3：2-Stage to Ping-Pong Pipeline。通过更高效的软件流水线掩盖访存延迟，实现计算与访存的重叠。

【这一部分内容待完善～】

FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling（2026）#

一种面向非对称硬件扩展的算法与 Kernel 流水线协同设计方法。

Algorithm#

Background#

相比 Hopper，Blackwell 的 Tensor Core 吞吐继续提升，但 shared memory 带宽、指数函数单元（MUFU）和普通 ALU 没有同比例增长，瓶颈从单纯的矩阵乘转移到 softmax、指数计算和 shared memory traffic。

Conditional Softmax Rescaling#

FA1/FA2 的 online softmax 在每个 block 都会根据新的 running max 缩放旧输出：

Oj=emj−1−mjOj−1+eSj−mjVj.O_j=e^{m_{j-1}-m_j}O_{j-1}+e^{S_j-m_j}V_j.Oj=emj−1−mjOj−1+eSj−mjVj.

FA4 观察到，只有当新的最大值明显变大时，旧输出才必须立即 rescale。因此它区分真实候选最大值 mjm_jmj 和当前用于缩放的基准 mˉj−1\bar{m}_{j-1}mˉj−1，只在：

mj−mˉj−1>τm_j-\bar{m}_{j-1}>\taumj−mˉj−1>τ

时才更新缩放基准并执行缩放：

mˉj←mj,Oj=emˉj−1−mˉjOj−1+eSj−mˉjVj.\bar{m}_j\leftarrow m_j,\qquad O_j=e^{\bar{m}_{j-1}-\bar{m}_j}O_{j-1}+e^{S_j-\bar{m}_j}V_j.mˉj←mj,Oj=emˉj−1−mˉjOj−1+eSj−mˉjVj.

否则不更新用于缩放的 mˉ\bar{m}mˉ，也不对旧输出做 rescale：

mˉj←mˉj−1,Oj=Oj−1+eSj−mˉj−1Vj.\bar{m}_j\leftarrow \bar{m}_{j-1},\qquad O_j=O_{j-1}+e^{S_j-\bar{m}_{j-1}}V_j.mˉj←mˉj−1,Oj=Oj−1+eSj−mˉj−1Vj.

阈值 τ\tauτ 控制可以容忍的 slack。论文中典型取 τ=log⁡2(256)=8\tau=\log_2(256)=8τ=log2(256)=8，对应“缩放因子不超过 256256256 时先不 rescale”。如果用自然指数 exe^xex 的写法理解，对应阈值就是 ln⁡(256)\ln(256)ln(256)。最后再用最终的 true normalizer 统一归一化，减少中间反复缩放带来的 non-matmul 开销。

Software-Emulated Exponential#

Blackwell 上 Tensor Core 吞吐增长很快，但指数函数单元（MUFU）没有同比例增长，softmax 里的 exp 会成为瓶颈。FA4 因此用 FMA 单元分担一部分指数计算。

softmax 数学上使用 exe^xex，但实现中可以先转成 base-2 exponent：

ex=2xlog⁡2e.e^x=2^{x\log_2 e}.ex=2xlog2e.

令：

y=xlog⁡2e,yint=⌊y⌋,yfrac=y−yint∈[0,1).y=x\log_2 e,\qquad y_{\text{int}}=\lfloor y\rfloor,\qquad y_{\text{frac}}=y-y_{\text{int}}\in[0,1).y=xlog2e,yint=⌊y⌋,yfrac=y−yint∈[0,1).

则：

ex=2y=2yint⋅2yfrac.e^x =2^y =2^{y_{\text{int}}}\cdot 2^{y_{\text{frac}}}.ex=2y=2yint⋅2yfrac.

整数部分 2yint2^{y_{\text{int}}}2yint 可以通过浮点数 exponent bits 处理；小数部分 2yfrac2^{y_{\text{frac}}}2yfrac 用低阶多项式近似：

2yfrac≈Pn(yfrac)=∑i=0npiyfraci.2^{y_{\text{frac}}} \approx P_n(y_{\text{frac}}) =\sum_{i=0}^{n}p_i y_{\text{frac}}^i.2yfrac≈Pn(yfrac)=i=0∑npiyfraci.

这不是把所有 exp 都替换成多项式，而是 partial emulation：一部分走 FMA 多项式近似，另一部分仍走 MUFU，避免寄存器压力和额外指令开销抵消收益。

其他优化#

【这一部分内容待完善～】

Source Code#

【这一部分内容待完善～】

参考资料#

KV Cache 与优化技术

Wed, 19 Aug 2026 01:00:00 GMT

Datawhale

大家好，我是芯缘，是 Datawhale 社区发起的 2026 年 8 月“llm-algo-leetcode 推理优化方向”组队学习活动的运营助教。本文记录了我学习 Task 1-2 KV Cache 的笔记。

KV Cache 的逻辑数据量为：

KV Cache 字节数=2×L×B×S×Hkv×D×e,\text{KV Cache 字节数} = 2 \times L \times B \times S \times H_{kv} \times D \times e,KV Cache 字节数=2×L×B×S×Hkv×D×e,

SSS 表示缓存的序列长度，eee 表示每个缓存标量占用的字节数：FP32 为 444，FP16/BF16 为 222，FP8/INT8 为 111。

注意力架构（Attention Architectures）#

MHA（Multi-Head Attention）是基线架构，MQA（Multi-Query Attention）、GQA（Grouped-Query Attention）和MLA（Multi-Head Latent Attention）以不同方式减少每个 Token 需要存储的表示。

表 1. 不同注意力架构及其 KV Cache 表示。

架构	缓存表示	相对于 MHA 的大小
MHA	Hkv=HqH_{kv}=H_qHkv=Hq	111
MQA	Hkv=1H_{kv}=1Hkv=1	1/Hq1/H_q1/Hq
GQA	1<Hkv<Hq1 < H_{kv} < H_q1<Hkv<Hq	Hkv/HqH_{kv}/H_qHkv/Hq
MLA	压缩后的 KV 潜在状态和解耦的 RoPE Key	取决于模型

MHA、GQA、MQA 和 MLA 对比

图 1. 条纹部分是在推理期间缓存的内容。MLA 存储压缩后的潜在 KV 表示，而不是为每个 Head 展开 Key 和 Value。来源：DeepSeek-V2, Figure 3。

表 2. 每条序列的 FP16/BF16 KV Cache 逻辑大小。

架构	KV Head 数量	4,096 个 Token	131,072 个 Token
MHA	32	2.00 GiB	64.00 GiB
GQA	8	0.50 GiB	16.00 GiB
MQA	1	0.06 GiB	2.00 GiB

这些数值不包括内存分配器填充（Allocator Padding）、Block Table、内存对齐（Alignment）和量化元数据（Quantization Metadata）。

KV Cache 优化技术#

多头潜在注意力（Multi-Head Latent Attention，MLA）#

MLA 优化的是每个 Token 存储的表示。它使用低秩联合压缩（Low-Rank Joint Compression），缓存紧凑的潜在状态和较小的位置分量，而不是为每个 Head 展开 K/V 状态。这会直接减少缓存的逻辑数据量，但需要采用 MLA 的模型和兼容的 Kernel。

滑动窗口注意力（Sliding-Window Attention）#

滑动窗口注意力优化的是保留的 Token 数量。每一层只保留固定大小的近期窗口，因此窗口填满后，缓存便不再增长。代价是被淘汰的 Token 无法再被直接关注。

StreamingLLM#

StreamingLLM 优化的是连续流式场景下的缓存淘汰策略。它保留少量起始的注意力汇聚 Token（Attention Sink Token）和一个近期 Token 窗口，在限制缓存大小的同时，比普通滑动窗口更不容易快速破坏模型的稳定性。被淘汰的中间 Token 仍然无法使用。

选择性 Token 淘汰与 KV Cache 剪枝#

选择性淘汰优化的是保留哪些 Token。它不只保留固定窗口，而是根据注意力计算得到的重要性或其他选择标准，在固定缓存预算下保留有用的较早 Token，并丢弃不太重要的条目。这是一种近似方法，因为某个 Token 一旦被淘汰，即使之后变得重要也无法恢复。

PagedAttention#

PagedAttention 优化的是物理内存分配。它将 KV Cache 划分为固定大小的 Block，并把逻辑上连续的 Token Block 映射到物理上不连续的 GPU 内存块，从而减少内存碎片和过度预留。它能提高内存利用率，但不会减少 K/V 本身的逻辑数据量。

前缀缓存（Prefix Caching）#

前缀缓存优化的是跨请求复用。具有相同前缀的请求可以复用之前已经计算的 K/V Block，从而减少重复的预填充（Prefill）计算和重复的前缀存储。它不会降低新 Token 的解码成本。

KV Cache 量化（KV Cache Quantization）#

KV Cache 量化优化的是存储精度。低比特格式能够减少每个 K/V 元素占用的字节数，代价是需要量化元数据，可能造成精度损失，并会引入量化或反量化开销。

CPU 卸载（CPU Offloading）#

CPU 卸载优化的是内存放置位置。它将一部分 KV Cache 从 GPU 显存移动到 CPU DRAM，在不改变缓存逻辑总量的情况下减少 GPU 驻留内存。主要代价是主机与设备之间的数据传输，以及由此可能带来的延迟。

后续#

上述每个优化方向都值得单独展开。以后将分别介绍 MLA、滑动窗口注意力、StreamingLLM、选择性 Token 淘汰与剪枝、PagedAttention、前缀缓存、KV Cache 量化和 CPU 卸载，包括各自的工作机制、实现约束和权衡。

参考资料#

GPU 架构与内存

Wed, 19 Aug 2026 00:00:00 GMT

Datawhale

大家好，我是芯缘，是 Datawhale 社区发起的 2026 年 8 月“llm-algo-leetcode 推理优化方向”组队学习活动的运营助教。本文记录了我学习 Task 1-1 GPU Hardware 的笔记。

GPU 架构#

GPU 组织结构#

CPU 与 GPU 组织结构

图 1. CPU 与 GPU 的整体组织结构及其内存连接。

一个 GPU 包含许多流式多处理器（Streaming Multiprocessor，SM）。每个 SM 都拥有自己的 CUDA Core、Tensor Core、寄存器文件（Register File）和共享内存（Shared Memory）/L1，而所有 SM 共享 L2 和 GPU DRAM。相比之下，CPU Core 共享容量较大的末级缓存（Last-Level Cache）和系统 DRAM。

一个 Kernel 会启动由多个线程块（Thread Block）组成的网格（Grid）。每个 Block 在一个 SM 上运行，并被划分为由 32 个线程组成的线程束（Warp），再由 SM 调度执行。

Blackwell SM 架构

图 2. Blackwell SM 的组织结构，包括 CUDA Core、Tensor Core、共享内存、TMA 和 TMEM。

Blackwell SM 新增了张量内存（Tensor Memory，TMEM），用于存放 Tensor Core 的累加器数据。张量内存加速器（Tensor Memory Accelerator，TMA）可以在全局内存（Global Memory）与共享内存之间搬运数据块（Tile），复制过程本身不占用 CUDA Core。

V100、A100、H100 与 B200#

下列表格采用数据中心 SXM 版本：V100 SXM2 32 GB、A100 80 GB SXM、H100 SXM 80 GB 和 B200 SXM。

表 1. 数据中心 GPU 规格。

规格	V100	A100	H100	B200
架构	Volta	Ampere	Hopper	Blackwell
发布年份	2017	2020	2022	2024
计算能力（Compute Capability）	7.0	8.0	9.0	10.0
SM 数量	80	108	132	148
Tensor Core 数量	640	432	528	592
L2 缓存	6 MB	40 MB	50 MB	126 MB
GPU 显存	32 GB HBM2	80 GB HBM2e	80 GB HBM3	180 GB HBM3e
显存带宽	900 GB/s	2.039 TB/s	3.35 TB/s	7.7 TB/s
最大功耗	300 W	400 W	700 W	最高 1,000 W

表 2. 各代 GPU 的主要架构变化。

GPU	主要架构变化
V100	首次引入第一代 Tensor Core，用于混合精度（Mixed Precision）矩阵运算。
A100	新增 TF32、BF16、结构化稀疏（Structured Sparsity）、多实例 GPU（Multi-Instance GPU，MIG）、异步复制（Asynchronous Copy），并大幅增大 L2 缓存。
H100	新增 FP8、Transformer Engine、张量内存加速器（TMA）、线程块集群（Thread-Block Cluster）和分布式共享内存。
B200	将两个受光刻掩模尺寸限制的裸片（Die）作为一个 GPU 使用，并新增 FP6、FP4 和第二代 Transformer Engine。

表 3. 不同精度下的稠密理论峰值算力。

精度	V100	A100	H100	B200
FP32（CUDA Core）	15.7	19.5	67	75
TF32（Tensor Core）	—	156	495	1,125
FP16（Tensor Core）	125	312	989	2,250
BF16（Tensor Core）	—	312	989	2,250
FP8（Tensor Core）	—	—	1,979	4,500
FP6（Tensor Core）	—	—	—	4,500
FP4（Tensor Core）	—	—	—	9,000

所有数值均为以 TFLOPS 计量的稠密理论峰值。FP32 使用 CUDA Core，其余各行使用 Tensor Core。

Tensor Core#

CUDA Core 为单个线程执行标量运算，而 Tensor Core 执行矩阵乘加（Matrix Multiply-Accumulate，MMA）：

D=AB+C.D = AB + C.D=AB+C.

Tensor Core 结合专用矩阵硬件、低精度输入以及通常更高精度的累加，来加速通用矩阵乘法（General Matrix Multiplication，GEMM）、卷积和 Transformer 层。只有当矩阵形状、数据布局、精度和数据搬运方式都与受支持的 Kernel 相匹配时，才能达到其峰值吞吐量。

内存与存储#

常用术语#

表 4. 容易混淆的缩写与名称。

术语	含义
L1/L2/L3（Level 1/2/3 Cache，一级/二级/三级缓存）	表示缓存层级；数字越小，通常距离计算单元越近、容量越小、速度越快。表 1 中的 GPU 拥有每个 SM 独立的 L1 和所有 SM 共享的 L2，但没有通用 L3。
共享内存（Shared Memory）	由程序员管理的片上内存，供一个 Block 内的线程共享。
显存（Video RAM，VRAM）	与 GPU 相连的内存，通常采用 GDDR 或 HBM；VRAM 描述的是用途角色，而不是存储单元技术。
固态硬盘（Solid-State Drive，SSD）	持久化存储设备，通常基于 NAND Flash。
机械硬盘（Hard Disk Drive，HDD）	使用旋转盘片的持久化磁存储设备。
NVMe（NVM Express）	一种通常承载于 PCIe 之上的 SSD 通信协议；NVM 指非易失性存储器（Nonvolatile Memory）。
PCIe（Peripheral Component Interconnect Express）	GPU、NVMe SSD 及其他设备使用的高速互连。
SATA（Serial ATA）	一种 HDD/SSD 接口；ATA 指 Advanced Technology Attachment。

内存与存储分类#

表 5. 常见内存与存储类别。数值范围为代表性数据，适用范围以各单元格中的说明为准。

大类	系列	常见形式与典型用途	带宽	容量
随机存取存储器（Random-Access Memory，RAM）易失性	静态随机存取存储器（Static RAM，SRAM）	GPU 共享内存/L1 和 L2 缓存	共享内存：128 B/cycle/SM L2：取决于产品	共享内存：96–228 KiB/SM L2：6–126 MB/GPU
动态随机存取存储器（Dynamic RAM，DRAM）	DDR SDRAM（Double Data Rate）——台式机/服务器内存	25–70 GB/s/channel	8 GB–4 TB/系统
LPDDR SDRAM（Low-Power DDR）——手机/笔记本电脑/Grace CPU 内存	50–500 GB/s/系统	6–480 GB/系统
GDDR SDRAM（Graphics DDR）——GPU 显存	0.3–1.8 TB/s	8–96 GB/GPU
HBM（High Bandwidth Memory，高带宽内存）——数据中心 GPU 显存	0.9–7.7 TB/s	32–180 GB/GPU
半导体非易失性存储器（Semiconductor Nonvolatile Memory，NVM）	只读存储器（Read-Only Memory，ROM）	固定式 ROM——启动代码/固件	取决于产品	KiB–GiB
EEPROM 与 Flash	EEPROM（Electrically Erasable Programmable ROM）——少量配置数据	取决于接口	Bytes–MiB
NOR Flash——BIOS/UEFI 与固件	读取速度 ≤400 MB/s	MiB–GiB
NAND Flash——SSD/手机/U 盘/SD 卡的存储介质	取决于接口和控制器	Gb–Tb/封装
存储设备（Storage Devices）非易失性	固态存储（Solid-State Storage）	SSD——SATA/NVMe 块存储，通常基于 NAND	SATA：0.5 GB/s NVMe：3–14 GB/s	0.25–122.88 TB/块
磁存储（Magnetic Storage）	HDD——机械式大容量存储	0.12–0.3 GB/s	1–36 TB/块

GPU 内存名称及其物理实现#

表 6. 无法仅从名称直接判断物理实现的 GPU 内存。

架构名称	名称所描述的对象	物理实现
GPU 寄存器（Register）	由 GPU 指令直接使用的每线程操作数	专用片上寄存器文件；NVIDIA 并未将其归类为普通 CUDA 共享内存
张量内存（Tensor Memory，TMEM）	Blackwell 中由第五代 Tensor Core 直接使用的累加器存储	与寄存器和共享内存分离的专用片上内存；其存储单元技术尚未公开说明
CUDA 局部内存（Local Memory）	每个线程私有的逻辑地址空间，用于寄存器溢出和部分自动变量	GPU 设备 DRAM，通过缓存层级访问；物理位置并不在相应线程附近
常量内存与纹理内存（Constant and Texture Memory）	具有专用缓存行为、以读取为主的 CUDA 地址空间	数据位于 GPU 设备 DRAM 中，通过专用片上缓存提供访问
分布式共享内存（Distributed Shared Memory）	访问同一线程块集群中其他 Block 的共享内存	由参与计算的各个 SM 上已有的片上共享内存资源构成
CUDA 统一内存（Unified Memory）	可由 CUDA 在不同处理器之间放置或迁移的托管虚拟内存分配	CPU DRAM 或 GPU DRAM，取决于数据当前的放置位置

GPU 内存层级（Memory Hierarchy）#

GPU 内存层级

图 3. 以 GPU 为中心的内存层级：越向下容量越大，越向上距离计算单元越近、速度越快。

高效的 Kernel 会尽量在寄存器和共享内存中复用数据，避免过早再次访问 HBM。

参考资料#

LeetCode | 双指针

Mon, 03 Aug 2026 15:38:00 GMT

分类：双指针 / 总结

讲解#

双指针题目的核心是：用两个指针描述当前正在处理的位置或范围，然后根据规则移动指针。

移动零的 left 和 right 一开始都在数组最左侧。right 不断向右扫描数组；当 right 看到非零数时，就把它交换到 left 指向的位置，然后 left 向右移动。最后数组前面是所有非零数，后面是所有 0。

盛最多水的容器的 left 一开始在最左侧，right 一开始在最右侧。两个指针夹住一个容器，计算当前面积后，移动较短的那一边；如果右边更短就移动 right，如果左边更短就移动 left。最后返回最大面积。

三数之和先排序并固定第一个数 nums[i]，然后让 left 从 i + 1 开始，right 从数组最右侧开始，在右侧区间里找另外两个数。如果三数之和偏小，就移动 left；如果三数之和偏大，就移动 right；如果等于 0，就记录答案并同时移动两个指针。最后返回所有不重复的三元组。

接雨水的 left 一开始在最左侧，right 一开始在最右侧，同时维护 left_max 和 right_max。如果左边高度更低，就结算 left 位置能接的水，然后移动 left；否则结算 right 位置能接的水，然后移动 right。最后返回总水量。

最长回文子串的 left 和 right 一开始都指向同一个中心位置（偶数长度回文则指向相邻两个字符）。只要 s[left] 和 s[right] 相等，就不断向两侧扩展，记录扩展出的回文串；每个中心都试过后，返回最长的回文子串。

双指针题目重要是：两个指针分别代表什么，什么时候移动左边，什么时候移动右边。

代码#

移动零#

1
from typing import List
2

3

4
class Solution:
5
    def moveZeroes(self, nums: List[int]) -> None:
6
        n = len(nums)
7
        left = right = 0
8
        while right < n:
9
            if nums[right]:
10
                nums[left], nums[right] = nums[right], nums[left]
11
                left += 1
12
            right += 1

盛最多水的容器#

1
from typing import List
2

3

4
class Solution:
5
    def maxArea(self, height: List[int]) -> int:
6
        left = 0
7
        right = len(height) - 1
8

9
        ans = 0
10
        while left < right:
11
            if height[right] < height[left]:
12
                h = height[right]
13
                area = (right - left) * h
14
                if area > ans:
15
                    ans = area
16
                right -= 1
17
                while left < right and height[right] < h:
18
                    right -= 1
19
            else:
20
                h = height[left]
21
                area = (right - left) * h
22
                if area > ans:
23
                    ans = area
24
                left += 1
25
                while left < right and height[left] < h:
26
                    left += 1
27

28
        return ans

展开收起

三数之和（双指针）#

1
from typing import List
2

3

4
class Solution:
5
    def threeSum(self, nums: List[int]) -> List[List[int]]:
6
        nums.sort()
7
        ans = []
8
        n = len(nums)
9

10
        for i in range(n - 2):
11
            if i > 0 and nums[i] == nums[i - 1]:
12
                continue
13
            if nums[i] > 0:
14
                break
15

16
            left = i + 1
17
            right = n - 1
18
            while left < right:
19
                total = nums[i] + nums[left] + nums[right]
20
                if total < 0:
21
                    left += 1
22
                elif total > 0:
23
                    right -= 1
24
                else:
25
                    ans.append([nums[i], nums[left], nums[right]])
26
                    left += 1
27
                    right -= 1
28
                    while left < right and nums[left] == nums[left - 1]:
29
                        left += 1
30
                    while left < right and nums[right] == nums[right + 1]:
31
                        right -= 1
32

33
        return ans

展开收起

接雨水#

1
from typing import List
2

3

4
class Solution:
5
    def trap(self, height: List[int]) -> int:
6
        ans = 0
7
        left = 0
8
        right = len(height) - 1
9
        left_max = 0
10
        right_max = 0
11

12
        while left < right:
13
            left_max = max(left_max, height[left])
14
            right_max = max(right_max, height[right])
15

16
            if height[left] < height[right]:
17
                ans += left_max - height[left]
18
                left += 1
19
            else:
20
                ans += right_max - height[right]
21
                right -= 1
22

23
        return ans

展开收起

最长回文子串（中心扩展）#

1
class Solution:
2
    def longestPalindrome(self, s: str) -> str:
3
        def expand(left: int, right: int) -> str:
4
            while left >= 0 and right < len(s) and s[left] == s[right]:
5
                left -= 1
6
                right += 1
7
            return s[left + 1:right]
8

9
        ans = ""
10
        for i in range(len(s)):
11
            odd = expand(i, i)
12
            even = expand(i, i + 1)
13
            ans = max(ans, odd, even, key=len)
14
        return ans

参考资料#

移动零
盛最多水的容器
三数之和
接雨水
最长回文子串
LeetCode 移动零：https://leetcode.com/problems/move-zeroes/
LeetCode 盛最多水的容器：https://leetcode.com/problems/container-with-most-water/
LeetCode 三数之和：https://leetcode.com/problems/3sum/
LeetCode 接雨水：https://leetcode.com/problems/trapping-rain-water/
LeetCode 最长回文子串：https://leetcode.com/problems/longest-palindromic-substring/
力扣移动零题解：https://leetcode.cn/problems/move-zeroes/solutions/489622/yi-dong-ling-by-leetcode-solution/
力扣盛最多水的容器题解：https://leetcode.cn/problems/container-with-most-water/solutions/11491/container-with-most-water-shuang-zhi-zhen-fa-yi-do/
力扣三数之和题解：https://leetcode.cn/problems/3sum/solutions/284681/san-shu-zhi-he-by-leetcode-solution/
力扣接雨水题解：https://leetcode.cn/problems/trapping-rain-water/solutions/692342/jie-yu-shui-by-leetcode-solution-tuvc/

LeetCode | 哈希表

Mon, 03 Aug 2026 15:37:00 GMT

分类：哈希表 / 总结

讲解#

哈希表题目的核心是：先建一张表，再去表里查东西。

两数之和把 已经遍历过的数字，以及它的下标放进字典作为哈希表,查target - num 在不在表里，最后返回两个数字对应的下标。

字母异位词分组把 排序后的字符串(key) -> 原字符串列表(value)放进字典作为哈希表,,查当前字符串排序后的 key在不在表里，最后返回字典里所有 value，也就是分好组的字符串列表。

最长连续序列把 nums 通过set(nums)去重后的所有数字放进集合作为哈希表,查先查 x - 1 在不在确定是不是序列起点，再不断查 x + 1、x + 2、…在不在表里，最后返回最长连续序列的长度。

哈希表题目重要是：我应该把什么变成 key，又要用这个 key 查什么。

代码#

两数之和#

1
class Solution:
2
    def twoSum(self, nums: list[int], target: int) -> list[int]:
3
        seen = {}
4

5
        for i, num in enumerate(nums):
6
            need = target - num
7
            if need in seen:
8
                return [seen[need], i]
9

10
            seen[num] = i
11
        return []

字母异位词分组（排序法）#

1
from collections import defaultdict
2
from typing import List
3

4

5
class Solution:
6
    def groupAnagrams(self, strs: List[str]) -> List[List[str]]:
7
        mp = defaultdict(list)
8

9
        for st in strs:
10
            key = "".join(sorted(st))
11
            mp[key].append(st)
12

13
        return list(mp.values())

最长连续序列#

1
from typing import List
2

3

4
class Solution:
5
    def longestConsecutive(self, nums: List[int]) -> int:
6
        st = set(nums)
7
        m = len(st)
8

9
        ans = 0
10
        for x in st:
11
            if x - 1 in st:
12
                continue
13

14
            y = x + 1
15
            while y in st:
16
                y += 1
17

18
            ans = max(ans, y - x)
19
            if ans * 2 >= m:
20
                break
21

22
        return ans

展开收起

参考资料#

两数之和
字母异位词分组
最长连续序列
LeetCode 两数之和：https://leetcode.com/problems/two-sum/
LeetCode 字母异位词分组：https://leetcode.com/problems/group-anagrams/
LeetCode 最长连续序列：https://leetcode.com/problems/longest-consecutive-sequence/
力扣两数之和题解：https://leetcode.cn/problems/two-sum/solutions/
力扣字母异位词分组题解：https://leetcode.cn/problems/group-anagrams/solutions/520469/zi-mu-yi-wei-ci-fen-zu-by-leetcode-solution/
力扣最长连续序列题解：https://leetcode.cn/problems/longest-consecutive-sequence/solutions/3005726/ha-xi-biao-on-zuo-fa-pythonjavacgojsrust-whop/

LeetCode | list、dict、set 有什么区别？

Fri, 31 Jul 2026 17:30:00 GMT

这个文档只整理几个最常见的容器：list、dict、set，以及字典分组时常用的 defaultdict(list)。

简单说：

1
list  -> 存一串东西，关心顺序，可以重复
2
dict  -> 存 key -> value 的对应关系
3
set   -> 存有没有出现过，不重复

`list`：列表 `[]`#

[] 表示列表，也就是 list。

列表适合存一串有顺序的数据，可以有重复元素。

常见操作：

1
arr.append(x)      # 在末尾添加 x
2
arr.pop()          # 删除并返回最后一个元素
3
arr.remove(x)      # 删除第一个等于 x 的元素
4
arr[i]             # 访问第 i 个元素

当你需要保存一组有顺序的数据，或者需要把一组结果收集起来时，用列表。

`dict`：字典 `{}`#

{} 通常表示字典，也就是 dict。

字典保存的是 key -> value 的对应关系。

1
mp = {}
2
mp["aet"] = ["eat", "tea", "ate"]

这里：

1
key   = "aet"
2
value = ["eat", "tea", "ate"]

常见操作：

1
mp[key] = value     # 修改
2
mp[key]             # 根据 key 取 value
3
key in mp           # 判断 key 是否存在
4
mp.values()         # 取出所有 value
5
del mp[key]         # 删除这个 key

当你想通过一个标识快速找到对应内容时，用字典。

比如在字母异位词分组里，可以用排序后的字符串 "aet" 找到对应的一组异位词。

注意：空的 {} 是字典，不是集合。

mp.values()#

mp.values() 取出来的是字典里的所有 value。比如：

1
mp = {
2
    "aet": ["eat", "tea", "ate"],
3
    "ant": ["tan", "nat"],
4
    "abt": ["bat"],
5
}
6

7
mp.values()

结果可以理解成：

1
[
2
    ["eat", "tea", "ate"],
3
    ["tan", "nat"],
4
    ["bat"]
5
]

如果题目要返回普通列表，通常再套一层 list：

1
list(mp.values())

字典分组：`defaultdict(list)`#

普通字典在 key 不存在时，不能直接这样写：

1
mp[key].append(st)

因为 mp[key] 还没有对应的列表，会报错。

普通字典需要先判断：

1
if key not in mp:
2
    mp[key] = []
3
mp[key].append(st)

defaultdict(list) 可以少写这个判断：

1
from collections import defaultdict
2

3
mp = defaultdict(list)
4
mp[key].append(st)

它可以理解成一个会自动创建空列表的字典。如果某个 key 还不存在，它会自动给这个 key 准备一个空列表。

当你要做“分组”时，经常可以用它。比如把同一类字符串放进同一个列表里。

`set`：集合 `set()`#

set() 表示集合。

集合只关心“有没有这个元素”，不保存重复元素，也不关心顺序。

常见操作：

1
s.add(x)         # 添加 x
2
s.remove(x)      # 删除 x，如果 x 不存在会报错
3
s.discard(x)     # 删除 x，如果 x 不存在也不报错
4
x in s           # 判断 x 是否存在

当你只关心某个东西有没有出现过时，用集合。

比如无重复字符的最长子串里，可以用集合记录当前窗口里有哪些字符。

小结#

1
list  -> 有顺序的一串数据
2
dict  -> 通过 key 快速找到 value
3
set   -> 快速判断某个东西有没有出现过

LeetCode | 三数之和

Mon, 03 Aug 2026 03:32:00 GMT

分类：双指针 / 数组 / 排序

题目#

给定一个整数数组 nums，找出所有和为 0 的三元组。

结果中不能包含重复的三元组。

示例：

1
输入：nums = [-1,0,1,2,-1,-4]
2
输出：[[-1,-1,2],[-1,0,1]]

讲解#

笔者的失败尝试#

一开始可以从两数之和的思路出发：固定第一个数 nums[first]，然后把问题变成在后面的数字里找两个数。

也就是：

1
nums[first] + nums[second] + nums[third] = 0

固定 nums[first] 后，就变成：

1
nums[second] + nums[third] = -nums[first]

这个思路本身是合理的，因为它确实把三数之和转成了两数之和的形式。

问题出在具体查找第三个数的写法：

1
target_third in nums[second + 1:]

这里的 nums[second + 1:] 不是在原列表上直接查找，而是先复制出一个新列表。切片复制需要 O(n)。

然后 target_third in 新列表 又是在列表里线性查找，也需要 O(n)。

所以这个方法整体接近：

1
外层 first        -> O(n)
2
内层 second       -> O(n)
3
切片和列表查找    -> O(n)
4
总时间复杂度      -> O(n^3)

虽然代码里用了 set 保存答案来去重，但它只是用来保存结果，不是用来做两数之和里的快速查找。所以这不是一个真正高效的哈希表解法。

这个方法在 LeetCode 上测试会超出时间限制。

基础思路#

暴力做法是枚举三个下标 i、j、k，检查 nums[i] + nums[j] + nums[k] 是否等于 0。这样有三层循环，时间复杂度是 O(n^3)。

这道题可以先排序，再用双指针。

排序之后，数组从小到大排列。我们先固定第一个数 nums[i]，然后在它右边找另外两个数。

此时问题变成：在 nums[i+1:] 里找两个数，让它们的和等于 -nums[i]

这就可以用双指针：

1
left  = i + 1
2
right = n - 1

如果三个数的和小于 0，说明当前和太小，需要让它变大，所以 left 往右移。

如果三个数的和大于 0，说明当前和太大，需要让它变小，所以 right 往左移。

如果刚好等于 0，就记录答案，然后两个指针都向中间移动。

这题还有一个关键点：去重。

固定第一个数时，如果 nums[i] == nums[i - 1]，说明这个第一个数之前已经用过，直接跳过。

找到一个答案之后，left 和 right 也要跳过相同的数字，否则会得到重复三元组。

1
left += 1
2
right -= 1
3
while left < right and nums[left] == nums[left - 1]:
4
    left += 1
5
while left < right and nums[right] == nums[right + 1]:
6
    right -= 1

剪枝#

第一个剪枝是：

1
if nums[i] > 0:
2
    break

数组已经排序，如果第一个数都大于 0，后面的数只会更大，不可能再凑出和为 0 的三元组。

笔者还尝试过另一个剪枝。固定 nums[i] 后，需要找两个数，让它们的和等于：

1
target = -nums[i]

第一版写法是先除以 2：

1
half = target / 2
2
if nums[left] > half or nums[right] < half:
3
    break

这个判断的意思是：

如果 nums[left] > half，说明当前最小的左边数已经太大了。因为 nums[right] >= nums[left]，所以两个数相加只会更大，不可能等于 target。
如果 nums[right] < half，说明当前最大的右边数也太小了。因为 nums[left] <= nums[right]，所以两个数相加只会更小，也不可能等于 target。

这个剪枝在数学上是成立的，但是 / 2 会引入浮点数计算，以及整数和浮点数之间的比较，开销更大。

后来把它改成了乘以 2 的版本：

1
target = -nums[i]
2
if nums[left] * 2 > target or nums[right] * 2 < target:
3
    break

这样可以避免浮点数，全部用整数比较。

但是加上这个剪枝以后，实际耗时还是增加了。原因是它需要在每次 while 循环里多做两次判断，而这个剪枝真正能提前结束循环的次数不够多，省下来的循环抵不过多出来的判断成本。

所以最终代码不加这个剪枝，只保留 nums[i] > 0 这个简单有效的剪枝。

复杂度#

暴力做法：
- 时间复杂度：O(n^3)
- 空间复杂度：O(1)
失败的方法（哈希表）：
- 时间复杂度：O(n^3)
- 空间复杂度：O(n)
排序 + 双指针：
- 时间复杂度：O(n^2)
- 空间复杂度：O(1)

失败的方法里有两层循环，里面的 nums[second + 1:] 和 in 查找又需要 O(n)，所以总时间复杂度接近 O(n^3)。ans 需要保存结果，切片也会临时创建新列表，所以空间复杂度按 O(n) 理解。

排序需要 O(n * log(n))，双指针部分是 O(n^2)，整体主要看 O(n^2)。

一些语法基础#

这道题里面有两个语法问题。

第一个是：

1
for i in range(n - 2):

这里不需要额外判断 n 有没有大于等于 3。

循环体会直接跳过，不会执行，也不会报错。

如果 n - 2 小于等于 0，比如 range(-1)、range(-2)，循环体会直接跳过，不会执行，也不会报错。

第二个是：

1
nums.sort()

sort() 是列表自己的方法，会直接修改原列表。

比如：

1
nums = [-1, 0, 1, 2, -1, -4]
2
nums.sort()

执行后，nums 自己会变成：

1
[-4, -1, -1, 0, 1, 2]

它和 sorted() 的区别是：

1
nums.sort()   -> 原地排序，修改 nums 本身，返回 None
2
sorted(nums)  -> 返回一个新的排序结果，不修改 nums 本身

在字母异位词分组里用的是 sorted(st)，因为那里要把字符串里的字符排序。字符串没有 sort() 方法，所以不能写 st.sort()。

代码#

失败的方法（哈希表）#

源码#

1
from typing import List
2

3

4
class Solution:
5
    def threeSum(self, nums: List[int]) -> List[List[int]]:
6
        length = len(nums)
7
        nums.sort()
8
        ans = set()
9

10
        for first in range(0, length):
11
            if nums[first] > 0:
12
                break
13
            target = 0 - nums[first]
14
            for second in range(first + 1, length):
15
                if nums[second] > target:
16
                    break
17
                target_third = target - nums[second]
18
                if target_third in nums[second + 1:]:
19
                    ans.add((nums[first], nums[second], target_third))
20

21
        return [list(item) for item in ans]

展开收起

双指针#

源码#

1
from typing import List
2

3

4
class Solution:
5
    def threeSum(self, nums: List[int]) -> List[List[int]]:
6
        nums.sort()
7
        ans = []
8
        n = len(nums)
9

10
        for i in range(n - 2):
11
            if i > 0 and nums[i] == nums[i - 1]:
12
                continue
13
            if nums[i] > 0:
14
                break
15

16
            left = i + 1
17
            right = n - 1
18
            while left < right:
19
                total = nums[i] + nums[left] + nums[right]
20
                if total < 0:
21
                    left += 1
22
                elif total > 0:
23
                    right -= 1
24
                else:
25
                    ans.append([nums[i], nums[left], nums[right]])
26
                    left += 1
27
                    right -= 1
28
                    while left < right and nums[left] == nums[left - 1]:
29
                        left += 1
30
                    while left < right and nums[right] == nums[right + 1]:
31
                        right -= 1
32

33
        return ans

展开收起

测试#

1
from src.python._015_Three_Sum import Solution
2

3

4
def normalize(ans):
5
    return sorted(sorted(x) for x in ans)
6

7

8
sol = Solution()
9

10
assert normalize(sol.threeSum([-1, 0, 1, 2, -1, -4])) == normalize([[-1, -1, 2], [-1, 0, 1]])
11
assert normalize(sol.threeSum([0, 1, 1])) == []
12
assert normalize(sol.threeSum([0, 0, 0])) == [[0, 0, 0]]
13
assert normalize(sol.threeSum([])) == []
14

15
print("PASS")

展开收起

参考资料#

LeetCode 原题：https://leetcode.com/problems/3sum/
力扣官方题解：https://leetcode.cn/problems/3sum/solutions/284681/san-shu-zhi-he-by-leetcode-solution/

LeetCode | 合并两个有序链表

Fri, 21 Aug 2026 04:00:00 GMT

分类：链表

题目#

将两个升序链表合并为一个新的升序链表并返回。新链表是通过拼接给定的两个链表的所有节点组成的。

示例：

1
输入：l1 = [1,2,4], l2 = [1,3,4]
2
输出：[1,1,2,3,4,4]

讲解#

两个链表已经有序，用两个指针 l1、l2 从头比较，每次把值更小的那个节点接到结果链表的末尾，谁的节点被取走谁就往后移。

关键点：

用 dummy 虚拟头节点串结果，最后返回 dummy.next，不用单独处理头节点（和两数相加一样的套路）。
循环条件是 l1 and l2，其中一个遍历完后，剩下的那段直接整体接上：cur.next = l1 if l1 else l2，因为剩下的本来就是有序的。

复杂度#

迭代法：
- 时间复杂度：O(m + n)
- 空间复杂度：O(1)

代码#

源码#

1
from typing import Optional
2

3

4
class ListNode:
5
    def __init__(self, val: int = 0, next: Optional['ListNode'] = None):
6
        self.val = val
7
        self.next = next
8

9

10
class Solution:
11
    def mergeTwoLists(self, l1: Optional[ListNode], l2: Optional[ListNode]) -> Optional[ListNode]:
12
        dummy = ListNode()
13
        cur = dummy
14

15
        while l1 and l2:
16
            if l1.val <= l2.val:
17
                cur.next = l1
18
                l1 = l1.next
19
            else:
20
                cur.next = l2
21
                l2 = l2.next
22
            cur = cur.next
23

24
        cur.next = l1 if l1 else l2
25
        return dummy.next

展开收起

测试#

1
def build_list(vals):
2
    dummy = ListNode()
3
    cur = dummy
4
    for v in vals:
5
        cur.next = ListNode(v)
6
        cur = cur.next
7
    return dummy.next
8

9

10
def to_list(head):
11
    res = []
12
    while head:
13
        res.append(head.val)
14
        head = head.next
15
    return res
16

17

18
sol = Solution()
19

20
assert to_list(sol.mergeTwoLists(build_list([1, 2, 4]), build_list([1, 3, 4]))) == [1, 1, 2, 3, 4, 4]
21
assert to_list(sol.mergeTwoLists(build_list([]), build_list([]))) == []
22
assert to_list(sol.mergeTwoLists(build_list([]), build_list([0]))) == [0]
23
assert to_list(sol.mergeTwoLists(build_list([1]), build_list([2]))) == [1, 2]
24

25
print("PASS")

展开收起

参考资料#

力扣原题：https://leetcode.cn/problems/merge-two-sorted-lists/

LeetCode | 两数相加

Fri, 21 Aug 2026 03:00:00 GMT

分类：链表

题目#

给你两个非空的链表，表示两个非负的整数。它们每位数字都是按照逆序的方式存储的，并且每个节点只能存储一位数字。

请你将两个数相加，并以相同形式返回一个表示和的链表。

你可以假设除了数字 0 之外，这两个数都不会以 0 开头。

示例：

1
输入：l1 = [9,9,9,9,9,9,9], l2 = [9,9,9,9]
2
输出：[8,9,9,9,0,0,0,1]

讲解#

两个链表都是逆序存数字，个位对齐排在链表头部，所以可以直接从头到尾同步遍历，模拟竖式相加。

维护一个 carry 进位。每一步把两个链表当前节点的值加上进位得到 total，当前位存 total % 10，进位存 total // 10。

关键点：

ListNode(total % 10) 创建一个新节点，值就是这一位的数字。cur.next = ... 把新节点挂到 cur 后面，再 cur = cur.next 移到新节点，下一轮继续串。
dummy 是占位的空节点，不存有效数字。循环里每次挂出的节点都跟在它后面，所以最后 return dummy.next 就是返回结果链表的头。有了 dummy，所有节点统一走 cur.next = 新节点，不用单独处理头节点。
两个链表长度可能不同，短的遍历完后按 0 处理。
while 条件带上 carry，这样最后一位进位 1 也能被处理，不会丢。

复杂度#

逐位相加：
- 时间复杂度：O(max(m, n))
- 空间复杂度：O(1)（返回值不计）

代码#

源码#

1
from typing import Optional
2

3

4
class ListNode:
5
    def __init__(self, val: int = 0, next: Optional['ListNode'] = None):
6
        self.val = val
7
        self.next = next
8

9

10
class Solution:
11
    def addTwoNumbers(self, l1: Optional[ListNode], l2: Optional[ListNode]) -> Optional[ListNode]:
12
        dummy = ListNode()
13
        cur = dummy
14
        carry = 0
15

16
        while l1 or l2 or carry:
17
            n1 = l1.val if l1 else 0
18
            n2 = l2.val if l2 else 0
19
            total = n1 + n2 + carry
20
            carry = total // 10
21
            cur.next = ListNode(total % 10)
22
            cur = cur.next
23
            if l1:
24
                l1 = l1.next
25
            if l2:
26
                l2 = l2.next
27

28
        return dummy.next

展开收起

测试#

1
def build_list(vals):
2
    dummy = ListNode()
3
    cur = dummy
4
    for v in vals:
5
        cur.next = ListNode(v)
6
        cur = cur.next
7
    return dummy.next
8

9

10
def to_list(head):
11
    res = []
12
    while head:
13
        res.append(head.val)
14
        head = head.next
15
    return res
16

17

18
sol = Solution()
19

20
assert to_list(sol.addTwoNumbers(build_list([2, 4, 3]), build_list([5, 6, 4]))) == [7, 0, 8]
21
assert to_list(sol.addTwoNumbers(build_list([0]), build_list([0]))) == [0]
22
assert to_list(sol.addTwoNumbers(build_list([9, 9, 9, 9, 9, 9, 9]), build_list([9, 9, 9, 9]))) == [8, 9, 9, 9, 0, 0, 0, 1]
23
assert to_list(sol.addTwoNumbers(build_list([5]), build_list([5]))) == [0, 1]
24

25
print("PASS")

展开收起

参考资料#

力扣原题：https://leetcode.cn/problems/add-two-numbers/

LeetCode | 最长回文子串

Fri, 21 Aug 2026 02:00:00 GMT

分类：双指针/字符串

题目#

给你一个字符串 s，找到 s 中最长的回文子串。

如果字符串的反序与原始字符串相同，则该字符串称为回文字符串。

示例 1：

1
输入：s = "babad"
2
输出："bab"
3
解释："aba" 同样是符合题意的答案。

示例 2：

1
输入：s = "cbbd"
2
输出："bb"

讲解#

回文串以中心向两侧对称，所以枚举每个位置当中心，向两边扩展比较。

关键点：回文中心有两种——奇数长度是单个字符 (i, i)，偶数长度是相邻两字符中间 (i, i + 1)。两种都要试。

expand(left, right) 从中心向两侧扩展，只要左右字符相等就继续外扩，直到不相等或越界。因为循环结束时左右指针已经多走了一步，所以切片取 s[left + 1:right]。

说明： 官方分类在多维动态规划。我这种题解是双指针。

复杂度#

中心扩展法：
- 时间复杂度：O(n^2)
- 空间复杂度：O(1)

代码#

源码#

1
class Solution:
2
    def longestPalindrome(self, s: str) -> str:
3
        def expand(left: int, right: int) -> str:
4
            while left >= 0 and right < len(s) and s[left] == s[right]:
5
                left -= 1
6
                right += 1
7
            return s[left + 1:right]
8

9
        ans = ""
10
        for i in range(len(s)):
11
            odd = expand(i, i)
12
            even = expand(i, i + 1)
13
            ans = max(ans, odd, even, key=len)
14
        return ans

测试#

1
sol = Solution()
2

3
assert sol.longestPalindrome("babad") in ("bab", "aba")
4
assert sol.longestPalindrome("cbbd") == "bb"
5
assert sol.longestPalindrome("a") == "a"
6
assert sol.longestPalindrome("ac") == "a"
7

8
print("PASS")

参考资料#

力扣原题：https://leetcode.cn/problems/longest-palindromic-substring/
题解参考：https://github.com/azl397985856/leetcode/blob/4183cbac/problems/5.longest-palindromic-substring.md

LeetCode | 反转链表

Wed, 12 Aug 2026 00:45:00 GMT

分类：链表

题目#

给你单链表的头节点 head，请你反转链表，并返回反转后的链表。

讲解#

反转链表的核心操作是让每个节点的 next 指向前一个节点。用 prev 记录前一个节点，curr 记录当前节点。

每一步：

先保存 curr.next 到 nxt（不然改了 curr.next 就找不到下一个了）。
curr.next = prev，完成当前节点的反转。
prev 和 curr 各往前移一步。

循环结束时 curr 为空，prev 就是新的头节点。

复杂度#

迭代法：
- 时间复杂度：O(n)
- 空间复杂度：O(1)

代码#

源码#

1
from typing import Optional
2

3

4
class ListNode:
5
    def __init__(self, val: int = 0, next: Optional['ListNode'] = None):
6
        self.val = val
7
        self.next = next
8

9

10
class Solution:
11
    def reverseList(self, head: Optional[ListNode]) -> Optional[ListNode]:
12
        prev = None
13
        curr = head
14

15
        while curr:
16
            nxt = curr.next
17
            curr.next = prev
18
            prev = curr
19
            curr = nxt
20

21
        return prev

展开收起

测试#

1
def build_list(vals):
2
    dummy = ListNode()
3
    cur = dummy
4
    for v in vals:
5
        cur.next = ListNode(v)
6
        cur = cur.next
7
    return dummy.next
8

9

10
def to_list(head):
11
    res = []
12
    while head:
13
        res.append(head.val)
14
        head = head.next
15
    return res
16

17

18
sol = Solution()
19

20
assert to_list(sol.reverseList(build_list([1, 2, 3, 4, 5]))) == [5, 4, 3, 2, 1]
21
assert to_list(sol.reverseList(build_list([1, 2]))) == [2, 1]
22
assert to_list(sol.reverseList(build_list([]))) == []
23
assert to_list(sol.reverseList(build_list([1]))) == [1]
24

25
print("PASS")

展开收起

参考资料#

力扣原题：https://leetcode.cn/problems/reverse-linked-list/
题解参考：https://github.com/doocs/leetcode/blob/main/solution/0200-0299/0206.Reverse%20Linked%20List/README.md

LeetCode | 有效的括号

Wed, 12 Aug 2026 00:40:00 GMT

分类：栈

题目#

给定一个只包括 '('，')'，'{'，'}'，'['，']' 的字符串 s，判断字符串是否有效。

有效字符串需满足：

左括号必须用相同类型的右括号闭合。
左括号必须以正确的顺序闭合。
每个右括号都有一个对应的相同类型的左括号。

示例 1：

1
输入：s = "(]"
2
输出：false

示例 2：

1
输入：s = "([])"
2
输出：true

讲解#

括号匹配是栈的经典应用。从左到右扫描字符串，用一个字典 pairs 存右括号到左括号的映射。

遇到左括号（不在 pairs 中）：入栈。遇到右括号：检查栈是否为空，或者栈顶是否匹配 pairs[ch]。不匹配直接返回 False。

提前判断：字符串长度是奇数直接返回 False，因为有效的括号一定是成对出现。

最后栈必须为空，才算全部匹配。

复杂度#

暴力做法：不断替换 "()"、"[]"、"{}" 直到字符串不再变化，时间复杂度 O(n²)。
栈：
- 时间复杂度：O(n)
- 空间复杂度：O(n)

代码#

源码#

1
class Solution:
2
    def isValid(self, s: str) -> bool:
3
        if len(s) % 2 == 1:
4
            return False
5

6
        pairs = {')': '(', ']': '[', '}': '{'}
7
        stack = []
8

9
        for ch in s:
10
            if ch not in pairs:
11
                stack.append(ch)
12
            elif not stack or stack.pop() != pairs[ch]:
13
                return False
14

15
        return not stack

展开收起

测试#

1
sol = Solution()
2

3
assert sol.isValid("()") == True
4
assert sol.isValid("()[]{}") == True
5
assert sol.isValid("(]") == False
6
assert sol.isValid("([])") == True
7
assert sol.isValid("([)]") == False
8
assert sol.isValid("{[]}") == True
9
assert sol.isValid("") == True
10
assert sol.isValid("((") == False
11

12
print("PASS")

参考资料#

力扣原题：https://leetcode.cn/problems/valid-parentheses/
题解参考：https://github.com/doocs/leetcode/blob/main/solution/0000-0099/0020.Valid%20Parentheses/README.md

LeetCode | 寻找两个正序数组的中位数

Sat, 08 Aug 2026 01:32:00 GMT

分类：二分查找 / 数组

写在前面#

这题应该是 Hot 100 里面最难的一道了。网上有传言戏称：“面试出这题感觉就是不想要你”。

面试官还请手下留情啊 (๑・́ω・̀๑)

题目#

给定两个大小分别为 m 和 n 的正序（从小到大）数组 nums1 和 nums2。请你找出并返回这两个正序数组的中位数。

算法的时间复杂度应该为 O(log (m+n)) 。

示例 1：

1
输入：nums1 = [1,3], nums2 = [2]
2
输出：2.00000
3
解释：合并数组 = [1,2,3] ，中位数 2

示例 2：

1
输入：nums1 = [1,2], nums2 = [3,4]
2
输出：2.50000
3
解释：合并数组 = [1,2,3,4] ，中位数 (2 + 3) / 2 = 2.5

讲解#

这道题最容易想到的方法是把两个数组合并，再找中位数。这样很好理解，但是需要看完所有元素，时间复杂度是 O(m+n)，不满足题目要求。

最优方法是二分查找。不过这里不是二分查找某个具体数字，而是二分查找一个正确的“切分位置”。

先不要看二分，先看中位数本身。

如果合并后的数组是 [1, 2, 3, 4, 5]，可以把它分成左边 [1, 2, 3] 和右边 [4, 5]。总长度是奇数，左边比右边多一个数，中位数就是左边最大的数，也就是 3。

如果合并后的数组是 [1, 2, 3, 4]，可以把它分成左边 [1, 2] 和右边 [3, 4]。总长度是偶数，中位数就是左边最大值和右边最小值的平均数，也就是 (2 + 3) / 2。

所以这道题真正要找的不是合并后的完整数组，而是一个切分方式。这个切分方式要满足两件事：左边的数量是对的，并且左边所有数都小于等于右边所有数。

假设在 nums1 里切 i 个数到左边，在 nums2 里切 j 个数到左边。为了保证左边数量正确，需要让 i + j = (m + n + 1) // 2。所以只要确定了 i，j 就能跟着确定：j = (m + n + 1) // 2 - i。

这就是为什么可以只在一个数组里二分：我们只需要二分 i，然后用总数量关系算出 j。

为了让二分更安全，代码先保证 nums1 是较短的数组。这样 i 的范围更小，j 也更不容易越界。

下面这张图展示的是一次合法切分：

两个正序数组的切分示意

每次切分后，只需要看切口附近四个数：nums1_left、nums1_right、nums2_left、nums2_right。

因为 nums1 和 nums2 本身都是有序的，所以 nums1 左边最大的数就是 nums1_left，右边最小的数就是 nums1_right；nums2 也是同理。

只要 nums1_left <= nums2_right，就说明 nums1 左边不会比 nums2 右边大。只要 nums2_left <= nums1_right，就说明 nums2 左边不会比 nums1 右边大。两个条件都满足，就代表整个左半边都小于等于整个右半边，这个切分就是正确的。

如果 nums1_left > nums2_right，说明 nums1 左边取多了，i 要往左移。否则说明 nums1 左边取少了，i 要往右移。

找到正确切分以后，中位数就只和切口旁边的数有关。如果总长度是奇数，中位数是 max(nums1_left, nums2_left)；如果总长度是偶数，中位数是 max(nums1_left, nums2_left) 和 min(nums1_right, nums2_right) 的平均数。

这里还用了 float("-inf") 和 float("inf") 处理边界。比如 i == 0 时，nums1 左边没有数，就把 nums1_left 看成负无穷；i == m 时，nums1 右边没有数，就把 nums1_right 看成正无穷。这样代码就不用单独为“切在最左边”或“切在最右边”写很多判断。

复杂度#

合并数组：
- 时间复杂度：O(m+n)
- 空间复杂度：O(m+n)
二分切分：
- 时间复杂度：O(log(min(m,n)))
- 空间复杂度：O(1)

代码#

源码#

1
from typing import List
2

3

4
class Solution:
5
    def findMedianSortedArrays(self, nums1: List[int], nums2: List[int]) -> float:
6
        if len(nums1) > len(nums2):
7
            return self.findMedianSortedArrays(nums2, nums1)
8

9
        m = len(nums1)
10
        n = len(nums2)
11
        total_left = (m + n + 1) // 2
12

13
        left = 0
14
        right = m
15

16
        while left <= right:
17
            i = (left + right) // 2
18
            j = total_left - i
19

20
            nums1_left = float("-inf") if i == 0 else nums1[i - 1]
21
            nums1_right = float("inf") if i == m else nums1[i]
22
            nums2_left = float("-inf") if j == 0 else nums2[j - 1]
23
            nums2_right = float("inf") if j == n else nums2[j]
24

25
            if nums1_left <= nums2_right and nums2_left <= nums1_right:
26
                if (m + n) % 2 == 1:
27
                    return float(max(nums1_left, nums2_left))
28
                return (max(nums1_left, nums2_left) + min(nums1_right, nums2_right)) / 2
29

30
            if nums1_left > nums2_right:
31
                right = i - 1
32
            else:
33
                left = i + 1
34

35
        return 0.0

展开收起

测试#

1
from src.python._004_Median_of_Two_Sorted_Arrays import Solution
2

3

4
sol = Solution()
5

6
assert sol.findMedianSortedArrays([1, 3], [2]) == 2.0
7
assert sol.findMedianSortedArrays([1, 2], [3, 4]) == 2.5
8
assert sol.findMedianSortedArrays([], [1]) == 1.0
9
assert sol.findMedianSortedArrays([2], []) == 2.0
10
assert sol.findMedianSortedArrays([0, 0], [0, 0]) == 0.0
11

12
print("PASS")

参考资料#

力扣原题：https://leetcode.cn/problems/median-of-two-sorted-arrays/
力扣题解页：https://leetcode.cn/problems/median-of-two-sorted-arrays/solutions/
二分切分讲解参考：https://liweiwei1419.github.io/leetcode-solution-blog/leetcode-problemset/binary-search/0004-median-of-two-sorted-arrays.html

LeetCode | 寻找旋转排序数组中的最小值

Sat, 08 Aug 2026 01:31:00 GMT

分类：二分查找 / 数组

题目#

已知一个长度为 n 的数组，预先按照升序排列，经由 1 到 n 次旋转后，得到输入数组。例如，原数组 nums = [0,1,2,4,5,6,7] 在变化后可能得到：

若旋转 4 次，则可以得到 [4,5,6,7,0,1,2]
若旋转 7 次，则可以得到 [0,1,2,4,5,6,7]

注意，数组 [a[0], a[1], a[2], ..., a[n-1]] 旋转一次的结果为数组 [a[n-1], a[0], a[1], a[2], ..., a[n-2]] 。

给你一个元素值互不相同的数组 nums ，它原来是一个升序排列的数组，并按上述情形进行了多次旋转。请你找出并返回数组中的最小元素。

你必须设计一个时间复杂度为 O(log n) 的算法解决此问题。

示例：

1
输入：nums = [3,4,5,1,2]
2
输出：1

讲解#

这题和搜索旋转排序数组很像，都是在旋转后的升序数组里做二分。

区别是：33 要找 target 的位置，这题只要找最小值。

最小值一定在旋转断点那里。比如 [4,5,6,7,0,1,2] 里，0 就是断点后的第一个数。

判断方向时，可以比较 nums[mid] 和 nums[right]。

如果 nums[mid] > nums[right]，说明 mid 在左边较大的那段里，最小值一定在 mid 右边，所以 left = mid + 1。

如果 nums[mid] <= nums[right]，说明从 mid 到 right 这一段是有序的，最小值可能就是 mid，也可能在 mid 左边，所以 right = mid。

循环结束时，left == right，这个位置就是最小值的位置。

复杂度#

暴力做法：
- 时间复杂度：O(n)
- 空间复杂度：O(1)
二分查找：
- 时间复杂度：O(log(n))
- 空间复杂度：O(1)

代码#

源码#

1
from typing import List
2

3

4
class Solution:
5
    def findMin(self, nums: List[int]) -> int:
6
        left = 0
7
        right = len(nums) - 1
8

9
        while left < right:
10
            mid = (left + right) // 2
11
            if nums[mid] > nums[right]:
12
                left = mid + 1
13
            else:
14
                right = mid
15

16
        return nums[left]

展开收起

测试#

1
from src.python._153_Find_Minimum_in_Rotated_Sorted_Array import Solution
2

3

4
sol = Solution()
5

6
assert sol.findMin([3, 4, 5, 1, 2]) == 1
7
assert sol.findMin([4, 5, 6, 7, 0, 1, 2]) == 0
8
assert sol.findMin([11, 13, 15, 17]) == 11
9
assert sol.findMin([1]) == 1
10

11
print("PASS")

参考资料#

力扣原题：https://leetcode.cn/problems/find-minimum-in-rotated-sorted-array/
力扣题解页：https://leetcode.cn/problems/find-minimum-in-rotated-sorted-array/solutions/

LeetCode | 搜索旋转排序数组

Sat, 08 Aug 2026 01:30:00 GMT

分类：二分查找 / 数组

题目#

整数数组 nums 按升序排列，数组中的值互不相同。

在传递给函数之前，nums 在预先未知的某个下标 k（0 <= k < nums.length）上进行了旋转，使数组变为 [nums[k], nums[k+1], ..., nums[n-1], nums[0], nums[1], ..., nums[k-1]]（下标从 0 开始计数）。例如， [0,1,2,4,5,6,7] 在下标 3 处经旋转后可能变为 [4,5,6,7,0,1,2]。

给你旋转后的数组 nums 和一个整数 target ，如果 nums 中存在这个目标值 target ，则返回它的下标，否则返回 -1 。

你必须设计一个时间复杂度为 O(log n) 的算法解决此问题。

示例 1：

1
输入：nums = [4,5,6,7,0,1,2], target = 0
2
输出：4

示例 2：

1
输入：nums = [4,5,6,7,0,1,2], target = 3
2
输出：-1

讲解#

这题和普通二分查找很像，但是数组被旋转过，整体不再完全升序。

关键点是：虽然整体不是升序，但每次看 left、mid、right 时，左右两半里面一定至少有一半是升序的。

如果 nums[left] <= nums[mid]，说明左半边 [left, mid] 是升序的。此时判断 target 是否落在 nums[left] 到 nums[mid] 之间：如果在，就去左半边找；否则去右半边找。

如果左半边不是升序的，那右半边 [mid, right] 一定是升序的。此时判断 target 是否落在 nums[mid] 到 nums[right] 之间：如果在，就去右半边找；否则去左半边找。

所以这题不是直接比较 nums[mid] 和 target 后决定左右，而是先判断哪一半有序，再判断 target 在不在有序的那一半里。

复杂度#

暴力做法：
- 时间复杂度：O(n)
- 空间复杂度：O(1)
二分查找：
- 时间复杂度：O(log(n))
- 空间复杂度：O(1)

代码#

源码#

1
from typing import List
2

3

4
class Solution:
5
    def search(self, nums: List[int], target: int) -> int:
6
        left = 0
7
        right = len(nums) - 1
8

9
        while left <= right:
10
            mid = (left + right) // 2
11
            if nums[mid] == target:
12
                return mid
13

14
            if nums[left] <= nums[mid]:
15
                if nums[left] <= target < nums[mid]:
16
                    right = mid - 1
17
                else:
18
                    left = mid + 1
19
            else:
20
                if nums[mid] < target <= nums[right]:
21
                    left = mid + 1
22
                else:
23
                    right = mid - 1
24

25
        return -1

展开收起

测试#

1
from src.python._033_Search_in_Rotated_Sorted_Array import Solution
2

3

4
sol = Solution()
5

6
assert sol.search([4, 5, 6, 7, 0, 1, 2], 0) == 4
7
assert sol.search([4, 5, 6, 7, 0, 1, 2], 3) == -1
8
assert sol.search([1], 0) == -1
9
assert sol.search([1], 1) == 0
10

11
print("PASS")

参考资料#

力扣原题：https://leetcode.cn/problems/search-in-rotated-sorted-array/
力扣题解页：https://leetcode.cn/problems/search-in-rotated-sorted-array/solutions/

LeetCode | 在排序数组中查找元素的第一个和最后一个位置

Fri, 07 Aug 2026 23:20:00 GMT

分类：二分查找 / 数组

题目#

给你一个按照非递减顺序排列的整数数组 nums，和一个目标值 target。请你找出给定目标值在数组中的开始位置和结束位置。

如果数组中不存在目标值 target，返回 [-1, -1]。

你必须设计并实现时间复杂度为 O(log n) 的算法解决此问题。

示例 1：

1
输入：nums = [5,7,7,8,8,10], target = 8
2
输出：[3,4]

示例 2：

1
输入：nums = [5,7,7,8,8,10], target = 6
2
输出：[-1,-1]

讲解#

这道题要找的是 target 出现的最左边位置和最右边位置。因为数组已经按照非递减顺序排列，所以相同的数字一定是连续挨在一起的。

这题可以直接复用搜索插入位置里的思路：写一个 lower_bound 函数，找第一个大于等于某个值的位置，然后调用两次。

第一次查找：第一个 >= target 的位置。这就是 target 应该出现的最左边位置。

如果这个位置越界，或者这个位置的值不是 target，说明数组里没有 target，直接返回 [-1, -1]。这里要注意：如果 start == len(nums)，那么 nums[start] 会越界。但是这段代码不会报错，因为 Python 的 or 会短路计算。左边 start == len(nums) 如果已经是 True，右边的 nums[start] != target 就不会再执行。所以这个判断是安全的。

第二次查找：第一个 >= target + 1 的位置。这个位置前面的一个位置，就是 target 出现的最后一个位置。

复杂度#

暴力做法：
- 时间复杂度：O(n)
- 空间复杂度：O(1)
二分查找：
- 时间复杂度：O(log(n))
- 空间复杂度：O(1)

代码#

源码#

1
from typing import List
2

3

4
class Solution:
5
    def searchRange(self, nums: List[int], target: int) -> List[int]:
6
        def lower_bound(target):
7
            left = 0
8
            right = len(nums)
9

10
            while left < right:
11
                mid = (left + right) // 2
12
                if nums[mid] < target:
13
                    left = mid + 1
14
                else:
15
                    right = mid
16

17
            return left
18

19
        start = lower_bound(target)
20
        if start == len(nums) or nums[start] != target:
21
            return [-1, -1]
22

23
        end = lower_bound(target + 1) - 1
24
        return [start, end]

展开收起

测试#

1
from src.python._034_Find_First_and_Last_Position_of_Element_in_Sorted_Array import Solution
2

3

4
sol = Solution()
5

6
assert sol.searchRange([5, 7, 7, 8, 8, 10], 8) == [3, 4]
7
assert sol.searchRange([5, 7, 7, 8, 8, 10], 6) == [-1, -1]
8
assert sol.searchRange([], 0) == [-1, -1]
9
assert sol.searchRange([1], 1) == [0, 0]
10

11
print("PASS")

参考资料#

力扣原题：https://leetcode.cn/problems/find-first-and-last-position-of-element-in-sorted-array/
力扣题解页：https://leetcode.cn/problems/find-first-and-last-position-of-element-in-sorted-array/solutions/
二分查找参考：https://en.wikipedia.org/wiki/Binary_search_algorithm

LeetCode | 搜索二维矩阵

Mon, 03 Aug 2026 17:00:00 GMT

分类：二分查找 / 数组 / 矩阵

题目#

给你一个满足下述两条属性的 m x n 整数矩阵：

每行中的整数从左到右按非严格递增顺序排列。
每行的第一个整数大于前一行的最后一个整数。

给你一个整数 target ，如果 target 在矩阵中，返回 true ；否则，返回 false 。

讲解#

二分法。

一维二分#

看成一维数组，在 [0, m * n) 这个范围里做二分，找第一个大于等于 target 的位置。

两次二分#

第二种方法是做两次二分。

第一次二分先找 target 可能在哪一行。

找到这一行之后，再对这一行做一次普通二分查找。

复杂度#

一维二分：
- 时间复杂度：O(log(m * n))
- 空间复杂度：O(1)
两次二分：
- 时间复杂度：O(log(m) + log(n))
- 空间复杂度：O(1)

因为 log(m) + log(n) = log(m * n)，所以两种方法的时间复杂度本质上是一样的。

代码#

一维二分#

源码#

1
from typing import List
2

3

4
class Solution:
5
    def searchMatrix(self, matrix: List[List[int]], target: int) -> bool:
6
        m = len(matrix)
7
        n = len(matrix[0])
8
        left = 0
9
        right = m * n
10

11
        while left < right:
12
            mid = (left + right) // 2
13
            x = matrix[mid // n][mid % n]
14
            if x < target:
15
                left = mid + 1
16
            else:
17
                right = mid
18

19
        return left < m * n and matrix[left // n][left % n] == target

展开收起

两次二分#

源码#

1
from typing import List
2

3

4
class Solution:
5
    def searchMatrix(self, matrix: List[List[int]], target: int) -> bool:
6
        m = len(matrix)
7
        n = len(matrix[0])
8

9
        top = 0
10
        bottom = m - 1
11
        while top <= bottom:
12
            row = (top + bottom) // 2
13
            if matrix[row][0] <= target <= matrix[row][n - 1]:
14
                break
15
            if matrix[row][0] > target:
16
                bottom = row - 1
17
            else:
18
                top = row + 1
19

20
        if top > bottom:
21
            return False
22

23
        left = 0
24
        right = n - 1
25
        while left <= right:
26
            mid = (left + right) // 2
27
            if matrix[row][mid] == target:
28
                return True
29
            if matrix[row][mid] < target:
30
                left = mid + 1
31
            else:
32
                right = mid - 1
33

34
        return False

展开收起

测试#

1
from src.python._074_Search_A_2D_Matrix import Solution
2

3

4
sol = Solution()
5

6
matrix = [[1, 3, 5, 7], [10, 11, 16, 20], [23, 30, 34, 60]]
7
assert sol.searchMatrix(matrix, 3) is True
8
assert sol.searchMatrix(matrix, 13) is False
9
assert sol.searchMatrix([[1]], 1) is True
10
assert sol.searchMatrix([[1]], 2) is False
11

12
print("PASS")

参考资料#

LeetCode 原题：https://leetcode.com/problems/search-a-2d-matrix/
AlgoMonster 一维二分题解：https://algo.monster/liteproblems/74
两次二分参考题解：https://blog.csdn.net/weixin_56980169/article/details/135609735

LeetCode | 搜索插入位置

Mon, 03 Aug 2026 16:09:00 GMT

分类：二分查找 / 数组

题目#

给定一个升序数组 nums 和一个目标值 target。

如果 target 在数组中，返回它的下标；如果不在数组中，返回它应该插入的位置。

示例：

1
输入：nums = [1,3,5,6], target = 5
2
输出：2

1
输入：nums = [1,3,5,6], target = 2
2
输出：1

讲解#

二分法，直接看代码。非常之清晰。二分法的范围通过双指针实现。

复杂度#

暴力做法：
- 时间复杂度：O(n)
- 空间复杂度：O(1)
二分查找：
- 时间复杂度：O(log(n))
- 空间复杂度：O(1)

代码#

源码#

1
from typing import List
2

3

4
class Solution:
5
    def searchInsert(self, nums: List[int], target: int) -> int:
6
        left = 0
7
        right = len(nums)
8

9
        while left < right:
10
            mid = (left + right) // 2
11
            if nums[mid] < target:
12
                left = mid + 1
13
            else:
14
                right = mid
15

16
        return left

展开收起

测试#

1
from src.python._035_Search_Insert_Position import Solution
2

3

4
sol = Solution()
5

6
assert sol.searchInsert([1, 3, 5, 6], 5) == 2
7
assert sol.searchInsert([1, 3, 5, 6], 2) == 1
8
assert sol.searchInsert([1, 3, 5, 6], 7) == 4
9
assert sol.searchInsert([1, 3, 5, 6], 0) == 0
10

11
print("PASS")

参考资料#

LeetCode | 二叉树的最大深度

Mon, 03 Aug 2026 15:36:00 GMT

分类：二叉树 / 深度优先搜索 / 递归

题目#

给定一棵二叉树的根节点 root，返回这棵树的最大深度。

最大深度指的是：从根节点到最远叶子节点经过的节点个数。

示例：

1
        3
2
       / \
3
      9   20
4
         /  \
5
        15   7
6

7
输入：root = [3,9,20,null,null,15,7]
8
输出：3

讲解#

如果对二叉树、根节点、左子树、右子树这些概念还不熟，可以先看二叉树的中序遍历里的基础概念。

当前树的最大深度 = max(左子树最大深度, 右子树最大深度) + 1

这个递归可以理解成：函数会不断往下找左子树和右子树，一直找到空节点。找到空节点时，说明这里已经没有节点了，所以返回 0。

然后结果会一层一层往上返回。每往上回到一个真实节点，就在左右子树最大深度的基础上 + 1，表示把当前这个节点也算进去。最后一直返回到最开始的根节点，就得到了整棵树的最大深度。

复杂度#

递归法：
- 时间复杂度：O(n)
- 空间复杂度：O(n)

每个节点都会被访问一次，所以时间复杂度是 O(n)。递归会使用系统调用栈，最坏情况下树退化成链表，递归深度是 n，所以空间复杂度是 O(n)。

代码#

源码#

1
class Solution:
2
    def maxDepth(self, root: "TreeNode") -> int:
3
        if not root:
4
            return 0
5

6
        left_height = self.maxDepth(root.left)
7
        right_height = self.maxDepth(root.right)
8
        return max(left_height, right_height) + 1

测试#

1
from src.python._104_Maximum_Depth_of_Binary_Tree import Solution
2

3

4
class TreeNode:
5
    def __init__(self, val=0, left=None, right=None):
6
        self.val = val
7
        self.left = left
8
        self.right = right
9

10

11
sol = Solution()
12

13
root = TreeNode(3)
14
root.left = TreeNode(9)
15
root.right = TreeNode(20)
16
root.right.left = TreeNode(15)
17
root.right.right = TreeNode(7)
18
assert sol.maxDepth(root) == 3
19

20
assert sol.maxDepth(None) == 0
21
assert sol.maxDepth(TreeNode(1)) == 1
22

23
root = TreeNode(1)
24
root.right = TreeNode(2)
25
assert sol.maxDepth(root) == 2
26

27
print("PASS")

展开收起

参考资料#

LeetCode 原题：https://leetcode.com/problems/maximum-depth-of-binary-tree/
力扣官方题解：https://leetcode.cn/problems/maximum-depth-of-binary-tree/solutions/349250/er-cha-shu-de-zui-da-shen-du-by-leetcode-solution/

LeetCode | 接雨水

Mon, 03 Aug 2026 06:59:00 GMT

分类：双指针 / 数组

题目#

给定一个非负整数数组 height，每个数字表示一个柱子的高度。

下雨之后，柱子之间可能会积水。求最多可以接多少水。

接雨水示例

示例：

1
输入：height = [0,1,0,2,1,0,1,3,2,1,2,1]
2
输出：6

讲解#

建议：先自己想一下“每个位置为什么能接水”，再看解析。否则双指针这一步会比较抽象。

这道题不要先想着“一整块水”怎么算，而是先看每一个位置能接多少水。

一个位置能接多少水，取决于它左右两边的挡板：

具体来说：

1
当前位置能接的水 = min(左边最高高度, 右边最高高度) - 当前高度

如果这个值小于 0，就说明这个位置接不了水。

比如某个位置高度是 1，它左边最高是 3，右边最高是 2，那么它最多只能接到高度 2，因为右边挡板更矮。

所以这个位置能接：

1
min(3, 2) - 1 = 1

暴力做法是：对每个位置，都向左找一次最高柱子，再向右找一次最高柱子，然后计算这个位置能接多少水。这样可以做出来，但是每个位置都要左右扫描，时间复杂度是 O(n^2)。

双指针的目标就是避免重复扫描。

我们让 left 从左往右走，right 从右往左走，同时维护：

1
left_max  = left 左侧已经见过的最高柱子
2
right_max = right 右侧已经见过的最高柱子

关键问题是：什么时候可以确定某个位置能接多少水？

当 height[left] < height[right] 时，可以先结算 left 这个位置。

原因是：右边此时至少有 height[right] 这根柱子，而且它比 height[left] 高。也就是说，left 右侧不是完全没有挡板，右边挡板已经够用了。

所以 left 位置能接多少水，主要取决于左边最高的柱子 left_max。

此时可以计算：

1
left_max - height[left]

然后让 left 往右移动。

反过来，当 height[left] >= height[right] 时，可以先结算 right 这个位置。

原因是：左边此时至少有 height[left] 这根柱子，而且它不低于 height[right]。所以 right 左侧挡板已经够用了。

此时 right 位置能接多少水，主要取决于右边最高的柱子 right_max。

此时可以计算：

1
right_max - height[right]

然后让 right 往左移动。

这就是双指针的核心：每次只结算“较矮的一边”。因为较矮的一边已经能确定自己的另一侧有挡板，不需要再等更远的位置。

复杂度#

暴力做法：
- 时间复杂度：O(n^2)
- 空间复杂度：O(1)
双指针做法：
- 时间复杂度：O(n)
- 空间复杂度：O(1)

每个指针只会向中间移动，不会回头，所以总移动次数最多是 n 次。

代码#

源码#

1
from typing import List
2

3

4
class Solution:
5
    def trap(self, height: List[int]) -> int:
6
        ans = 0
7
        left = 0
8
        right = len(height) - 1
9
        left_max = 0
10
        right_max = 0
11

12
        while left < right:
13
            left_max = max(left_max, height[left])
14
            right_max = max(right_max, height[right])
15

16
            if height[left] < height[right]:
17
                ans += left_max - height[left]
18
                left += 1
19
            else:
20
                ans += right_max - height[right]
21
                right -= 1
22

23
        return ans

展开收起

测试#

1
from src.python._042_Trapping_Rain_Water import Solution
2

3

4
sol = Solution()
5

6
assert sol.trap([0, 1, 0, 2, 1, 0, 1, 3, 2, 1, 2, 1]) == 6
7
assert sol.trap([4, 2, 0, 3, 2, 5]) == 9
8
assert sol.trap([1, 2, 3]) == 0
9
assert sol.trap([]) == 0
10

11
print("PASS")

参考资料#

LeetCode 原题：https://leetcode.com/problems/trapping-rain-water/
力扣官方题解：https://leetcode.cn/problems/trapping-rain-water/solutions/692342/jie-yu-shui-by-leetcode-solution-tuvc/

LeetCode | 最长连续序列

Sun, 02 Aug 2026 09:14:00 GMT

分类：哈希表 / 数组

题目#

给定一个未排序的整数数组 nums，找出数字连续的最长序列长度。

这个连续序列不要求在原数组里连续出现。

示例：

1
输入：nums = [100, 4, 200, 1, 3, 2]
2
输出：4
3
原因：最长连续序列是 [1, 2, 3, 4]

讲解#

暴力做法是：把每个数字都当作起点，然后不断查找 x + 1、x + 2、x + 3 是否存在。这样会重复计算很多次。

优化的关键是：只从连续序列的起点开始数。通过判断x-1在不在集合内来确定x是不是起点。

为了快速判断某个数字是否存在，先把数组转成集合：

1
st = set(nums)

后面还有一个剪枝：

1
if ans * 2 >= m:
2
    break

这里 m 是去重后的数字个数。如果当前已经找到的连续序列长度 ans 至少占了一半，那么答案不可能再变得更大。

复杂度#

暴力做法：
- 时间复杂度：O(n^2)
- 空间复杂度：O(n)
哈希集合优化：
- 时间复杂度：O(n)
- 空间复杂度：O(n)

代码#

源码#

1
from typing import List
2

3

4
class Solution:
5
    def longestConsecutive(self, nums: List[int]) -> int:
6
        st = set(nums)
7
        m = len(st)
8

9
        ans = 0
10
        for x in st:
11
            if x - 1 in st:
12
                continue
13

14
            y = x + 1
15
            while y in st:
16
                y += 1
17

18
            ans = max(ans, y - x)
19
            if ans * 2 >= m:
20
                break
21

22
        return ans

展开收起

测试#

1
from src.python._128_Longest_Consecutive_Sequence import Solution
2

3

4
sol = Solution()
5

6
assert sol.longestConsecutive([100, 4, 200, 1, 3, 2]) == 4
7
assert sol.longestConsecutive([0, 3, 7, 2, 5, 8, 4, 6, 0, 1]) == 9
8
assert sol.longestConsecutive([]) == 0
9
assert sol.longestConsecutive([1, 1, 1]) == 1
10

11
print("PASS")

参考资料#

LeetCode 原题：https://leetcode.com/problems/longest-consecutive-sequence/
力扣题解：https://leetcode.cn/problems/longest-consecutive-sequence/solutions/3005726/ha-xi-biao-on-zuo-fa-pythonjavacgojsrust-whop/

LeetCode | 盛最多水的容器

Sun, 02 Aug 2026 06:59:00 GMT

分类：双指针 / 数组

题目#

给定一个整数数组 height，数组中的每个数字表示一条竖线的高度。

任意选择两条竖线，它们和 x 轴可以组成一个容器。

求能够装下的最大水量。

盛最多水的容器示例

示例：

1
输入：height = [1,8,6,2,5,4,8,3,7]
2
输出：49

讲解#

容器能装多少水，取决于：

1
宽度 = 两条线之间的距离
2
高度 = 两条线中较短的那条
3
面积 = 宽度 * 高度

暴力做法是枚举所有两条线的组合，计算每一对的面积。这样可以得到答案，但是有两层循环，时间复杂度是 O(n^2)。

这道题可以用双指针。先让 left 指向最左边，right 指向最右边。

当前面积是：

1
(right - left) * min(height[left], height[right])

接下来应该移动较短的那一边。

原因是：如果固定较短的那条线，只移动较高的那条线，宽度一定变小，而容器高度仍然不会超过原来的短板，所以面积不可能变大。

笔者的小巧思： 代码在这个基础上又做了一步剪枝：移动短板之后，如果新的高度还比刚才的短板更矮，就继续跳过。因为宽度已经变小，高度还更矮，这种位置也不可能得到更大的面积。

复杂度#

暴力做法：
- 时间复杂度：O(n^2)
- 空间复杂度：O(1)
双指针做法：
- 时间复杂度：O(n)
- 空间复杂度：O(1)

每个指针只会向中间移动，不会回头，所以总移动次数最多是 n 次。

代码#

源码#

1
from typing import List
2

3

4
class Solution:
5
    def maxArea(self, height: List[int]) -> int:
6
        left = 0
7
        right = len(height) - 1
8

9
        ans = 0
10
        while left < right:
11
            if height[right] < height[left]:
12
                h = height[right]
13
                area = (right - left) * h
14
                if area > ans:
15
                    ans = area
16
                right -= 1
17
                while left < right and height[right] < h:
18
                    right -= 1
19
            else:
20
                h = height[left]
21
                area = (right - left) * h
22
                if area > ans:
23
                    ans = area
24
                left += 1
25
                while left < right and height[left] < h:
26
                    left += 1
27

28
        return ans

展开收起

测试#

1
from src.python._011_Container_With_Most_Water import Solution
2

3

4
sol = Solution()
5

6
assert sol.maxArea([1, 8, 6, 2, 5, 4, 8, 3, 7]) == 49
7
assert sol.maxArea([1, 1]) == 1
8
assert sol.maxArea([4, 3, 2, 1, 4]) == 16
9

10
print("PASS")

参考资料#

LeetCode 原题：https://leetcode.com/problems/container-with-most-water/
力扣官方题解：https://leetcode.cn/problems/container-with-most-water/solutions/11491/container-with-most-water-shuang-zhi-zhen-fa-yi-do/

LeetCode | 字母异位词分组

Fri, 31 Jul 2026 17:20:00 GMT

分类：哈希表 / 字符串 / 排序 / 计数

题目#

给定一个字符串数组 strs，把所有字母异位词分到同一组。

字母异位词指的是：两个字符串包含的字母完全相同，只是顺序不同。

示例：

1
输入：strs = ["eat", "tea", "tan", "ate", "nat", "bat"]
2
输出：[["bat"], ["nat", "tan"], ["ate", "eat", "tea"]]

返回结果中，每一组内部的顺序、组和组之间的顺序都不重要。

讲解#

这道题的关键是：同一组字母异位词需要有同一个标识。

排序法#

把字符串里的字符sort排序后，字母异位词会得到同一个结果。所以可以把排序后的字符串作为哈希表的 key。

具体做法：

准备一个字典 mp。
遍历每个字符串 st。
对 st 里的字符排序，得到 key。
把原字符串 st 放进 mp[key] 对应的列表里。

计数法#

字母异位词不一定要排序才能判断，也可以数每个字母出现了几次。

所以可以用一个长度为 26 的数组记录每个小写字母出现次数。这个数组转成元组后，也可以作为哈希表的 key。

计数法避免了排序，所以每个字符串只需要扫描一遍。

复杂度#

这里 n 是字符串数量，k 是单个字符串的最大长度。

排序法：
- 时间复杂度：O(n * k * log(k))
- 空间复杂度：O(n * k)
计数法：
- 时间复杂度：O(n * k)
- 空间复杂度：O(n * k)

排序法里的 log(k) 来自排序。对一个长度为 k 的字符串排序，常见排序算法的时间复杂度是 O(k * log(k))。

可以用“归并排序”来理解这个复杂度。假设有 k 个字符要排序：

1
第 1 层：把 k 个字符分成 2 组
2
第 2 层：继续分成 4 组
3
第 3 层：继续分成 8 组
4
...

每往下一层，问题规模大约减半，所以总层数大约是 log(k)。

分完以后还要一层一层合并。每一层合并时，所有字符加起来都会被处理一遍，也就是 k 次左右。

所以可以理解成：

1
每一层处理 k 个字符
2
一共有 log(k) 层
3
总时间复杂度 = O(k * log(k))

Python 的 sorted 底层不是简单的归并排序，而是 Timsort。不过从这道题分析复杂度时，可以按常见排序的 O(k * log(k)) 来理解。

计数法不排序，只扫描每个字符串里的字符。扫描一个长度为 k 的字符串是 O(k)，扫描 n 个字符串就是 O(n * k)。

一些语法基础#

sorted(st) 会把字符串 st 里的字符拿出来排序，返回的是列表，不是字符串。sorted(“eat”)结果是：[“a”, “e”, “t”]

join 是字符串的方法，格式是：连接符.join(字符串列表)，例如”-“.join([“a”, “e”, “t”])结果是：“a-e-t”

计数法里会用到 ord(ch) - ord("a")。ord(ch) 会得到字符对应的数字编号，所以减去 ord("a") 之后，就能把小写字母转成 0 到 25 的下标：

tuple(counts) 是把列表转成元组。字典的 key 必须是不可变的，列表可以修改，所以不能当 key；元组不能修改，所以可以当 key。计数法里要把 counts 转成 tuple 再放进字典。

代码#

排序法#

源码#

1
from collections import defaultdict
2
from typing import List
3

4

5
class Solution:
6
    def groupAnagrams(self, strs: List[str]) -> List[List[str]]:
7
        mp = defaultdict(list)
8

9
        for st in strs:
10
            key = "".join(sorted(st))
11
            mp[key].append(st)
12

13
        return list(mp.values())

计数法#

源码#

1
from collections import defaultdict
2
from typing import List
3

4

5
class Solution:
6
    def groupAnagrams(self, strs: List[str]) -> List[List[str]]:
7
        mp = defaultdict(list)
8

9
        for st in strs:
10
            counts = [0] * 26
11
            for ch in st:
12
                counts[ord(ch) - ord("a")] += 1
13
            mp[tuple(counts)].append(st)
14

15
        return list(mp.values())

展开收起

测试#

1
def normalize(groups):
2
    return sorted(sorted(group) for group in groups)
3

4

5
sol = Solution()
6

7
ans = sol.groupAnagrams(["eat", "tea", "tan", "ate", "nat", "bat"])
8
assert normalize(ans) == normalize([["bat"], ["nat", "tan"], ["ate", "eat", "tea"]])
9

10
assert normalize(sol.groupAnagrams([""])) == normalize([[""]])
11
assert normalize(sol.groupAnagrams(["a"])) == normalize([["a"]])
12

13
print("PASS")

参考资料#

LeetCode 原题：https://leetcode.com/problems/group-anagrams/
力扣官方题解：https://leetcode.cn/problems/group-anagrams/solutions/520469/zi-mu-yi-wei-ci-fen-zu-by-leetcode-solution/

LeetCode | 矩阵置零

Fri, 31 Jul 2026 01:50:00 GMT

分类：数组 / 矩阵

题目#

给定一个 m x n 的矩阵 matrix。如果某个元素为 0，就把它所在的整行和整列都设为 0。

要求直接修改原矩阵。

示例：

1
输入：
2
[
3
  [1, 1, 1],
4
  [1, 0, 1],
5
  [1, 1, 1]
6
]
7

8
输出：
9
[
10
  [1, 0, 1],
11
  [0, 0, 0],
12
  [1, 0, 1]
13
]

讲解#

这道题不能看到 0 就立刻把整行整列改成 0。

原因是：新改出来的 0 会影响后面的判断，导致本来不该被清零的位置也被清零。

所以要分两步：

先记录哪些行、哪些列需要变成 0。
再统一把这些行和列改成 0。

这里直接使用标记数组：

row_0[row] 表示第 row 行要不要变成 0
col_0[col] 表示第 col 列要不要变成 0

先扫描一遍矩阵。只要看到 matrix[row][col] == 0，就把 row_0[row] 和 col_0[col] 标记为 1。

再扫描一遍矩阵。如果当前位置的行被标记了，或者列被标记了，就把这个位置改成 0。

复杂度：

时间复杂度：O(mn)
空间复杂度：O(m + n)

吐槽： 官方题解还有两个标记变量和一个标记变量的写法，时间复杂度都是 O(mn)，空间复杂度可以做到 O(1)。它们会复用矩阵的第一行、第一列做标记，笔者觉得，可读性不如这个版本；实际开发里一般也不至于为了这点空间把代码写得那么绕，所以这里不写。

代码#

源码#

1
from typing import List
2

3

4
class Solution:
5
    def setZeroes(self, matrix: List[List[int]]) -> None:
6
        """
7
        Do not return anything, modify matrix in-place instead.
8
        """
9
        m = len(matrix)
10
        n = len(matrix[0])
11

12
        row_0 = [0] * m
13
        col_0 = [0] * n
14

15
        for row in range(0,m):
16
            for col in range (0,n):
17
                if matrix[row][col] == 0 :
18
                    row_0[row],col_0[col] = 1,1
19

20
        for row in range(0,m):
21
            for col in range (0,n):
22
                if (row_0[row] or col_0[col]) :
23
                    matrix[row][col] = 0

展开收起

测试#

1
from src.python._073_Set_Matrix_Zeroes import Solution
2

3

4
sol = Solution()
5

6
matrix = [[1, 1, 1], [1, 0, 1], [1, 1, 1]]
7
assert sol.setZeroes(matrix) is None
8
assert matrix == [[1, 0, 1], [0, 0, 0], [1, 0, 1]]
9

10
matrix = [[0, 1, 2, 0], [3, 4, 5, 2], [1, 3, 1, 5]]
11
sol.setZeroes(matrix)
12
assert matrix == [[0, 0, 0, 0], [0, 4, 5, 0], [0, 3, 1, 0]]
13

14
matrix = [[1, 0]]
15
sol.setZeroes(matrix)
16
assert matrix == [[0, 0]]
17

18
matrix = [[1], [0]]
19
sol.setZeroes(matrix)
20
assert matrix == [[0], [0]]
21

22
print("PASS")

展开收起

参考资料#

LeetCode 原题：https://leetcode.com/problems/set-matrix-zeroes/
力扣官方题解：https://leetcode.cn/problems/set-matrix-zeroes/solutions/669901/ju-zhen-zhi-ling-by-leetcode-solution-9ll7/

LeetCode | 无重复字符的最长子串

Fri, 31 Jul 2026 01:40:00 GMT

分类：哈希集合 / 滑动窗口 / 字符串

题目#

给定一个字符串 s，找出其中不含重复字符的最长子串长度。

这里的子串必须是连续的一段。

示例：

1
输入：s = "abcabcbb"
2
输出：3
3
原因：答案可以是 "abc"，长度为 3。"bca" 和 "cab" 也可以。

1
输入：s = "bbbbb"
2
输出：1
3
原因：答案是 "b"，长度为 1。

1
输入：s = "pwwkew"
2
输出：3
3
原因：答案是 "wke"，长度为 3。"pwke" 不是子串，因为它不是连续的一段。

讲解#

这道题要找的是“没有重复字符”的最长连续片段。

遇到下面几种信号时，经常可以考虑滑动窗口：

题目要求的是连续的一段，比如子串、子数组
需要找最长、最短、最多这类范围问题
右边界可以一步步往右扩
如果窗口不满足条件，左边界可以一步步往右缩

这道题正好符合这些信号：子串必须连续，目标是找最长；没有重复时可以继续往右扩，出现重复时就从左边缩小窗口。

做法是滑动窗口。窗口可以理解成字符串里当前正在看的连续片段。

我们维护一个窗口 [left, right]，并保证窗口里没有重复字符。right 从左到右扫描字符串，left 只在遇到重复字符时往右移动。

curr 是一个集合，用来记录当前窗口里有哪些字符。

当新字符 ch 已经在 curr 里时，说明窗口里已经有一个同样的字符了。这个重复字符不一定就在最左边，所以不能假设 s[left] 一定是它。

代码的做法是：从窗口最左边开始删，一个一个删，直到窗口里不再有 ch。

1
while ch in curr:
2
    curr.remove(s[left])
3
    left += 1

注意这里用的是 while。意思是：只要新字符 ch 还在窗口里，就继续缩小窗口。

以 s = "abcabcbb" 为例。当 right 走到第二个 a 时：

1
当前窗口: "abc"
2
curr = {a, b, c}
3
新字符 ch = "a"

此时 "a" 已经在窗口里。刚好最左边也是 "a"，删掉它以后窗口变成 "bc"，就可以把新的 "a" 加进来，窗口变成 "bca"。

但重复字符不总是在最左边。看 s = "abba"：

1
窗口先变成 "ab"
2
接下来遇到第二个 "b"

这时重复的是 "b"，但窗口最左边是 "a"。所以第一步先删 "a"，窗口变成 "b"；可是 "b" 还在窗口里，所以还要继续删 "b"。删到窗口里没有 "b" 以后，才能把新的 "b" 加进来。

所以这段代码不是因为“最左边一定是重复字符”才删最左边，而是因为窗口要保持连续。想让左边界右移，只能从最左边开始一个个移除，直到重复字符被移出窗口。

这个方法快的原因是：left 和 right 都只会从左往右走，不会回头。

复杂度：

滑动窗口：时间复杂度 O(n)，空间复杂度 O(n)

代码#

源码#

1
class Solution:
2
    def lengthOfLongestSubstring(self, s: str) -> int:
3
        curr = set()
4
        left = 0
5
        ans = 0
6

7
        for right,ch in enumerate(s):
8
            while ch in curr:
9
                curr.remove(s[left])
10
                left += 1
11
            curr.add(ch)
12
            ans = max(right - left + 1 , ans)
13

14
        return ans

测试#

1
from src.python._003_Longest_Substring_Without_Repeating_Characters import Solution
2

3

4
sol = Solution()
5

6
assert sol.lengthOfLongestSubstring("abcabcbb") == 3
7
assert sol.lengthOfLongestSubstring("bbbbb") == 1
8
assert sol.lengthOfLongestSubstring("pwwkew") == 3
9
assert sol.lengthOfLongestSubstring("") == 0
10
assert sol.lengthOfLongestSubstring(" ") == 1
11
assert sol.lengthOfLongestSubstring("abba") == 2
12

13
print("PASS")

参考资料#

LeetCode 原题：https://leetcode.com/problems/longest-substring-without-repeating-characters/
力扣题解页：https://leetcode.cn/problems/longest-substring-without-repeating-characters/solutions/

LeetCode | 二叉树的中序遍历

Fri, 31 Jul 2026 01:30:00 GMT

分类：二叉树 / 深度优先搜索 / 栈

题目#

给定一棵二叉树的根节点 root，返回它的中序遍历结果。

中序遍历的顺序是：左子树、当前节点、右子树。

示例：

1
        4
2
      /   \
3
     2     6
4
    /     / \
5
   1     5   7
6
输入：root = [4,2,6,1,null,5,7]
7
输出：[1,2,4,5,6,7]

其余例子略

讲解#

基础概念#

树是一种像倒过来的树一样的数据结构：从一个起点往下分叉。

节点就是树里的一个元素。每个节点通常会保存一个值，比如图里的 1、2、3。

根节点就是整棵树最上面的节点，也就是题目里的 root。如果 root = []，说明这棵树是空的。

二叉树表示每个节点最多有两个孩子：左孩子和右孩子。以某个节点的左孩子为起点继续往下看，就是它的左子树；以右孩子为起点继续往下看，就是它的右子树。

深度优先搜索就是先沿着一条路尽量往深处走，走不下去了再退回来换另一条路。二叉树的前序、中序、后序遍历都属于这种思路。

中序遍历专门指二叉树里的“左、根、右”：先访问左子树，再访问当前节点，最后访问右子树。

栈是LIFO。stack.append(x) 表示把 x 放进栈，stack.pop() 表示取出最后放进去的元素。

这道题的核心就是记住中序遍历的访问顺序：左、根、右。

递归法#

递归法最直观，但要先注意一点：代码里的 root 不一定永远是整棵树最上面的根节点，它更像是“当前正在处理的节点”。

如果左边或右边没有节点，对应的 root.left 或 root.right 就是 None。

递归函数 inorder(root) 做的事情是：

如果当前节点是空的，直接返回。
先递归处理当前节点的左子树。
左边处理完以后，把当前节点的值加入结果。
最后递归处理当前节点的右子树。

以 root = [4,2,6,1,null,5,7] 为例：

执行过程可以简单看成：

1
inorder(4)
2
  inorder(4.left) -> inorder(2)
3
    inorder(2.left) -> inorder(1)
4
      inorder(1.left) -> 空，返回
5
      加入 1
6
      inorder(1.right) -> 空，返回
7
    加入 2
8
    inorder(2.right) -> 空，返回
9
  加入 4
10
  inorder(4.right) -> inorder(6)
11
    inorder(6.left) -> inorder(5)
12
      inorder(5.left) -> 空，返回
13
      加入 5
14
      inorder(5.right) -> 空，返回
15
    加入 6
16
    inorder(6.right) -> inorder(7)
17
      inorder(7.left) -> 空，返回
18
      加入 7
19
      inorder(7.right) -> 空，返回

展开收起

所以结果是 [1, 2, 4, 5, 6, 7]。

迭代栈法#

递归本质上也在用系统调用栈。迭代法就是自己准备一个 stack 来模拟递归。

中序遍历是“左、根、右”，所以一个节点不能一看到就立刻加入结果，必须先把它左边的节点处理完。

栈里保存的就是这些“还不能访问，先等一下”的节点。

具体过程是：

当前节点不为空时，一直往左走，路过的节点都放进栈。
走到空节点，说明左边已经到底了。
从栈里 pop 一个节点，这个节点的左边已经处理完，现在可以加入结果。
然后转向这个节点的右子树，继续重复。

以 root = [4,2,6,1,null,5,7] 为例：

1
        4
2
      /   \
3
     2     6
4
    /     / \
5
   1     5   7

执行过程可以简单看成：

1
while root or stack:
2
            while root:
3
                stack.append(root)
4
                root = root.left
5
放入 4，往左到 2
6
放入 2，往左到 1
7
放入 1，往左到空
8

9
root = stack.pop()
10
res.append(root.val)
11
root = root.right
12
pop 1，加入 1，1 没有右边（跳出while root，继续执行上面）
13
pop 2，加入 2，2 没有右边
14

15
pop 4，加入 4，转向 4 的右边 6
16

17
放入 6，往左到 5
18
放入 5，往左到空
19
pop 5，加入 5，5 没有右边
20

21
pop 6，加入 6，转向 6 的右边 7
22
放入 7，往左到空
23
pop 7，加入 7，7 没有右边

展开收起

所以结果是 [1, 2, 4, 5, 6, 7]。

复杂度#

递归法：时间复杂度 O(n)，空间复杂度 O(n)
迭代栈法：时间复杂度 O(n)，空间复杂度 O(n)

代码#

递归法#

源码#

1
class Solution:
2
    def inorderTraversal(self, root: "TreeNode") -> list[int]:
3
        res = []
4

5
        def inorder(root):
6
            if not root:
7
                return
8
            inorder(root.left)
9
            res.append(root.val)
10
            inorder(root.right)
11

12
        inorder(root)
13
        return res

迭代栈法#

源码#

1
class Solution:
2
    def inorderTraversal(self, root: "TreeNode") -> list[int]:
3
        res = []
4
        stack = []
5

6
        while root or stack:
7
            while root:
8
                stack.append(root)
9
                root = root.left
10

11
            root = stack.pop()
12
            res.append(root.val)
13
            root = root.right
14

15
        return res

展开收起

测试#

1
class TreeNode:
2
    def __init__(self, val=0, left=None, right=None):
3
        self.val = val
4
        self.left = left
5
        self.right = right
6

7

8
sol = Solution()
9

10
root = TreeNode(1)
11
root.right = TreeNode(2)
12
root.right.left = TreeNode(3)
13
assert sol.inorderTraversal(root) == [1, 3, 2]
14

15
root = TreeNode(4)
16
root.left = TreeNode(2)
17
root.right = TreeNode(6)
18
root.left.left = TreeNode(1)
19
root.right.left = TreeNode(5)
20
root.right.right = TreeNode(7)
21
assert sol.inorderTraversal(root) == [1, 2, 4, 5, 6, 7]
22

23
root = TreeNode(1)
24
root.left = TreeNode(2)
25
root.right = TreeNode(3)
26
root.left.left = TreeNode(4)
27
root.left.right = TreeNode(5)
28
root.left.right.left = TreeNode(6)
29
root.left.right.right = TreeNode(7)
30
root.right.right = TreeNode(8)
31
root.right.right.left = TreeNode(9)
32
assert sol.inorderTraversal(root) == [4, 2, 6, 5, 7, 1, 3, 9, 8]
33

34
assert sol.inorderTraversal(None) == []
35
assert sol.inorderTraversal(TreeNode(1)) == [1]
36

37
print("PASS")

展开收起

参考资料#

LeetCode 原题与配图：https://leetcode.com/problems/binary-tree-inorder-traversal/
力扣题解页：https://leetcode.cn/problems/binary-tree-inorder-traversal/solutions/

LeetCode | 相交链表

Fri, 31 Jul 2026 01:20:00 GMT

分类：链表 / 哈希集合 / 双指针

题目#

给定两个单链表的头节点 headA 和 headB，找出两个链表相交的起始节点。如果两个链表不相交，返回 None。

这里的“相交”不是节点值相等，而是两个链表从某个节点开始共用同一段节点。

平台会根据 intersectVal、listA、listB、skipA、skipB 构造链表，但这些不是函数的直接输入。函数真正接收的是两个头节点 headA 和 headB。

要求函数返回后，两个链表仍然保持原来的结构；题目保证整个链式结构中没有环。

示例：

相交链表示例 1

1
intersectVal = 8
2
listA = [4,1,8,4,5]
3
listB = [5,6,1,8,4,5]
4
skipA = 2
5
skipB = 3
6
输出：相交节点 8

其余例子略

讲解#

这道题要找的是“同一个节点对象”，不是值相同的节点。

哈希集合#

先走一遍链表 A，把 A 里面出现过的节点都放进一个集合；再走一遍链表 B，第一个已经在集合里的节点，就是相交的起始节点。

set() 是 Python 里的集合。它只关心“某个东西在不在里面”，不保存下标，也不保存对应关系。这里 visited = set() 表示创建一个空集合，visited.add(cur) 表示把当前节点放进去，cur in visited 表示判断当前节点以前是否出现过。

它和两数之和里的哈希表很像，都是利用哈希让查找平均为 O(1)。区别是：两数之和用字典保存“数字 -> 下标”，这里用集合只保存“这个节点出现过”。

长度差对齐#

先分别算出两条链表的长度 lenA 和 lenB。如果长链表比短链表多 s 个节点，就先让长链表走 s 步。这样两个指针到尾部的距离就一样了。

之后两个指针一起往后走：如果它们指向同一个节点，这个节点就是相交的起始节点；如果一直没有相同节点，最后会同时走到 None，说明不相交。

复杂度#

哈希集合法：时间复杂度 O(m + n)，空间复杂度 O(m)
长度差对齐法：时间复杂度 O(m + n)，空间复杂度 O(1)

代码#

哈希集合法#

源码#

1
class Solution:
2
    def getIntersectionNode(self, headA: ListNode, headB: ListNode) -> Optional[ListNode]:
3
        _hash =set();
4

5
        pointerA = headA
6
        while (pointerA):
7
            _hash.add(pointerA)
8
            pointerA = pointerA.next
9

10
        pointerB = headB
11
        while (pointerB):
12
            if pointerB in _hash:
13
                return pointerB
14
            pointerB = pointerB.next
15

16
        return None

展开收起

长度差对齐法#

源码#

1
class Solution:
2
    def getIntersectionNode(self, headA: "ListNode", headB: "ListNode") -> "ListNode":
3
        lenA = self.getLength(headA)
4
        lenB = self.getLength(headB)
5

6
        if lenA >= lenB:
7
            long, short = headA, headB
8
        else:
9
            long, short = headB, headA
10

11
        for _ in range(abs(lenA - lenB)):
12
            long = long.next
13

14
        while long != short:
15
            long = long.next
16
            short = short.next
17

18
        return long
19

20
    def getLength(self, head: "ListNode") -> int:
21
        length = 0
22
        while head:
23
            length += 1
24
            head = head.next
25
        return length

展开收起

测试#

1
class ListNode:
2
    def __init__(self, x):
3
        self.val = x
4
        self.next = None
5

6

7
sol = Solution()
8

9
common = ListNode(8)
10
common.next = ListNode(4)
11
common.next.next = ListNode(5)
12

13
headA = ListNode(4)
14
headA.next = ListNode(1)
15
headA.next.next = common
16

17
headB = ListNode(5)
18
headB.next = ListNode(6)
19
headB.next.next = ListNode(1)
20
headB.next.next.next = common
21

22
assert sol.getIntersectionNode(headA, headB) is common
23

24
headA = ListNode(2)
25
headA.next = ListNode(6)
26
headA.next.next = ListNode(4)
27

28
headB = ListNode(1)
29
headB.next = ListNode(5)
30

31
assert sol.getIntersectionNode(headA, headB) is None
32
assert sol.getIntersectionNode(None, ListNode(1)) is None
33

34
print("PASS")

展开收起

参考资料#

LeetCode 原题与配图：https://leetcode.com/problems/intersection-of-two-linked-lists/
力扣官方题解：https://leetcode.cn/problems/intersection-of-two-linked-lists/solutions/811625/xiang-jiao-lian-biao-by-leetcode-solutio-a8jn/

LeetCode | 移动零

Fri, 31 Jul 2026 01:10:00 GMT

分类：双指针

题目#

给定一个整数数组 nums，需要把数组中的所有 0 移动到末尾，同时保持非零元素原来的相对顺序。

要求必须原地修改数组，不能复制一个新数组作为答案。

示例：

1
输入: nums = [0, 1, 0, 3, 12]
2
输出: [1, 3, 12, 0, 0]

1
输入: nums = [0]
2
输出: [0]

讲解#

这道题的关键信号是：原地修改数组，并且要保持非零元素的相对顺序。

目标数组最终会变成：

1
所有非零数 + 所有 0

所以我们只需要关心一件事：下一个非零数应该放在哪里。

双指针正好适合这种场景。

right 从左到右扫描数组，负责找非零数；

left 表示下一个非零数应该放到的位置。

移动零双指针过程

双指针只扫描一遍数组，时间复杂度是 O(n)，空间复杂度是 O(1)。

代码#

源码#

1
from typing import List
2
class Solution:
3
    def moveZeroes(self, nums: List[int]) -> None:
4
        left = 0
5
        for right in range(0,len(nums)):
6
            if(nums[right] != 0):
7
                nums[left],nums[right] = nums[right],nums[left]
8
                left+=1

测试#

1
from src.python._283_Move_Zeroes import Solution
2

3

4
sol = Solution()
5

6
nums = [0, 1, 0, 3, 12]
7
assert sol.moveZeroes(nums) is None
8
assert nums == [1, 3, 12, 0, 0]
9

10
nums = [0]
11
sol.moveZeroes(nums)
12
assert nums == [0]
13

14
nums = [1, 0, 1]
15
sol.moveZeroes(nums)
16
assert nums == [1, 1, 0]
17

18
print("PASS")

展开收起

参考资料#

LeetCode 原题：https://leetcode.com/problems/move-zeroes/
力扣官方题解：https://leetcode.cn/problems/move-zeroes/solutions/489622/yi-dong-ling-by-leetcode-solution/

LeetCode | 两数之和

Fri, 31 Jul 2026 01:00:00 GMT

分类：哈希表

题目#

给定一个整数数组 nums 和一个整数 target，需要找出数组中两个不同位置的元素，使它们的和等于 target，并返回这两个元素的下标。

题目保证每组输入恰好存在一个答案。同一个元素不能使用两次。返回的两个下标顺序通常不重要。

示例：

1
输入: nums = [2, 7, 11, 15], target = 9
2
输出: [0, 1]
3
原因: nums[0] + nums[1] = 2 + 7 = 9

1
输入: nums = [3, 2, 4], target = 6
2
输出: [1, 2]

1
输入: nums = [3, 3], target = 6
2
输出: [0, 1]

讲解#

目标是找两个数，它们相加等于 target。

最直接的做法是两层循环，把所有组合都试一遍。时间复杂度是 O(n^2)，空间复杂度是 O(1)。

更好的做法是用字典记录已经见过的数字。Python 的字典就是一种哈希表，可以理解成“快速查找表”，这里保存的是：

1
数字 -> 下标

当遍历到当前数字 num 时，只需要看字典里有没有 target - num。如果有，这两个数就正好组成答案。

哈希表会根据要查的数字快速定位位置，不是像列表那样从头到尾一个个比较。need in seen 查的是字典，平均查找复杂度是 O(1)，所以不会把整体复杂度变回 O(n^2)。

关键点是：先查需要的数，再记录当前数。这样不会把同一个元素用两次。

复杂度：

暴力做法：时间复杂度 O(n^2)，空间复杂度 O(1)
字典做法：时间复杂度 O(n)，空间复杂度 O(n)

代码#

源码#

1
class Solution:
2
    def twoSum(self, nums: List[int], target: int) -> List[int]:
3
        _hash = {}
4

5
        for i,num in enumerate(nums):
6
            find = target - num
7
            if( find in _hash ):
8
                return [_hash[find],i]
9
            _hash[num] = i
10
        return []

测试#

1
from src.python._001_Two_Sum import Solution
2

3

4
sol = Solution()
5
assert sol.twoSum([2, 7, 11, 15], 9) == [0, 1]
6
assert sol.twoSum([3, 2, 4], 6) == [1, 2]
7
assert sol.twoSum([3, 3], 6) == [0, 1]
8

9
print("PASS")

参考资料#

LeetCode 原题：https://leetcode.com/problems/two-sum/
力扣官方题解：https://leetcode.cn/problems/two-sum/solutions/434597/liang-shu-zhi-he-by-leetcode-solution/

LeetGPU | 02_Matrix_Multiplication

Sat, 25 Jul 2026 00:00:00 GMT

关注 LeetGPU repository，了解更多内容。

Problem Description#

LeetGPU Challenge

Write a program that multiplies two matrices of 32-bit floating point numbers on a GPU. Given matrix AAA of dimensions M×NM \times NM×N and matrix BBB of dimensions N×KN \times KN×K, compute the product matrix C=A×BC = A \times BC=A×B, which will have dimensions M×KM \times KM×K. All matrices are stored in row-major format.

Implementation Requirements#

Use only native features (external libraries are not permitted).
The solve function signature must remain unchanged.
The final result must be stored in matrix C.

Example 1#

Input:

Matrix AAA (2×22 \times 22×2):

[1.02.03.04.0]\begin{bmatrix} 1.0 & 2.0 \\ 3.0 & 4.0 \end{bmatrix}[1.03.02.04.0]

Matrix BBB (2×22 \times 22×2):

[5.06.07.08.0]\begin{bmatrix} 5.0 & 6.0 \\ 7.0 & 8.0 \end{bmatrix}[5.07.06.08.0]

Output:

Matrix CCC (2×22 \times 22×2):

[19.022.043.050.0]\begin{bmatrix} 19.0 & 22.0 \\ 43.0 & 50.0 \end{bmatrix}[19.043.022.050.0]

Example 2#

Input:

Matrix AAA (1×31 \times 31×3):

[1.02.03.0]\begin{bmatrix} 1.0 & 2.0 & 3.0 \end{bmatrix}[1.02.03.0]

Matrix BBB (3×13 \times 13×1):

[4.05.06.0]\begin{bmatrix} 4.0 \\ 5.0 \\ 6.0 \end{bmatrix}4.05.06.0

Output:

Matrix CCC (1×11 \times 11×1):

[32.0]\begin{bmatrix} 32.0 \end{bmatrix}[32.0]

Constraints#

1 <= M, N, K <= 8192.
Performance is measured with M = 8192, N = 6144, K = 4096.

CUDA#

Approach#

本题需要计算 A[M, N] 与 B[N, K] 的乘积，并将结果写入 C[M, K]。每个输出元素都是 A 的一行与 B 的一列所形成的内积（inner product）：

C[row,col]=∑inner=0N−1A[row,inner]×B[inner,col].C[\text{row},\text{col}] = \sum_{\text{inner}=0}^{N-1} A[\text{row},\text{inner}] \times B[\text{inner},\text{col}].C[row,col]=inner=0∑N−1A[row,inner]×B[inner,col].

由于三个矩阵均采用行主序（row-major）存储，它们的二维坐标可分别转换为以下线性地址：

1
A[row, inner] -> A[row * N + inner]
2
B[inner, col] -> B[inner * K + col]
3
C[row, col]   -> C[row * K + col]

不同 (row, col) 坐标对应的输出彼此独立，因此 CUDA 可以并行处理输出矩阵的两个维度；沿 N 维度进行的归约（reduction），则仍由单个线程（thread）或该线程负责的输出分块（tile）完成。

传统 CPU 实现通常使用三重嵌套循环：外侧两个循环遍历 M×KM \times KM×K 个输出位置，最内层循环完成长度为 NNN 的归约，整体算术复杂度为 O(MNK)O(MNK)O(MNK)。GPU 优化并不会减少这些必要的算术运算，而是重新安排循环的执行位置、同时计算的输出数量，以及数据从各级内存中被读取的频率。

矩阵乘法的内积视角

一个输出元素由 A 的一行和 B 的一列计算得到。

当前 Markdown 查看器不支持嵌入式视频。请直接打开[输出元素动画](/posts/assets/02_Matrix_Multiplication/video_01.mp4)。

这些 CUDA 文件展示了一条逐步演进的优化路径。所有版本都保持相同的 solve(...) 接口和数学结果，主要区别在于每个 thread 负责多少输出，以及数据如何在全局内存（global memory）、共享内存（shared memory）与寄存器（register）之间流动。

整个优化过程围绕两个相互关联的目标展开：

提高复用率。 从 global memory 读取的数据，应在被替换前尽可能参与多个输出元素的计算。共享内存分块（shared-memory tiling）可以在同一个 block 内复用数据，线程分块（thread tiling）则能进一步在 register 中复用数据。
让数据移动与计算重叠，或提高二者的执行效率。 向量化传输可以减少 load/store 指令开销；软件流水线和硬件流水线可以让相邻 reduction tile 的数据移动与计算交叠；如果能够接受较低的 TF32 输入精度，还可以使用 Tensor Core 取代标量 FP32 乘加指令。

Version	Main change	Work owned by one thread or warp
V0	Direct global-memory dot products	One output element per thread
V1 shared	Cooperative shared-memory tiling	One output element per thread
V1 1D	Register reuse along the row dimension	`16 x 1` outputs per thread
V1 2D	Register-level outer products	`8 x 4` outputs per thread
V2	`float4` cooperative transfers	`8 x 4` outputs per thread
V3	Register prefetch plus two shared-memory buffers	`8 x 4` outputs per thread
V4	Hardware-asynchronous `cp.async` transfers	`8 x 4` outputs per thread
V5	WMMA TF32 Tensor Core operations	One `16 x 16` output tile per warp

V0: Naive Global-Memory Kernel#

Approach#

V0 启动一个二维网格（grid），其中每个线程块（block）包含 16 x 16 个 thread。每个 thread 按如下方式确定自己负责的输出坐标：

1
row = blockIdx.y * blockDim.y + threadIdx.y
2
col = blockIdx.x * blockDim.x + threadIdx.x

grid 在两个输出维度上都使用向上取整除法计算大小。每个 thread 首先检查 row < M && col < K；只有坐标有效的 thread 才会计算完整的点积，并将结果写入 C 中对应的一个元素。

这种映射用二维 CUDA grid 取代了 CPU 实现最外侧的两个循环，但每个 thread 仍需串行完成整个内部归约。该实现结构直接，因而也很容易验证正确性。

V0 的主要弱点在于数据移动。分别计算 C[row, col0] 和 C[row, col1] 的两个 thread 都会遍历 A 的同一行，却各自重新加载这行数据；同样，计算同一输出列中不同行元素的 thread，也会反复读取 B 同一列中的值。硬件缓存（cache）或许能够捕获一部分复用机会，但该 kernel 并未显式地把这些数据保留在片上。因此，V0 虽然暴露了充足的输出并行性，却没有改善底层点积计算的数据局部性。

Solution#

1
#include <cuda_runtime.h>
2

3
__global__ void matrix_multiplication_v0(const float* A, const float* B, float* C,
4
                                         int M, int N, int K) {
5
    const int row = blockIdx.y * blockDim.y + threadIdx.y;
6
    const int col = blockIdx.x * blockDim.x + threadIdx.x;
7
    if (row >= M || col >= K) return;
8

9
    float sum = 0.0f;
10
    for (int inner = 0; inner < N; ++inner) {
11
        sum += A[row * N + inner] * B[inner * K + col];
12
    }
13
    C[row * K + col] = sum;
14
}
15

16
extern "C" void solve(const float* A, const float* B, float* C, int M, int N, int K) {
17
    const dim3 block(16, 16);
18
    const dim3 grid((K + block.x - 1) / block.x, (M + block.y - 1) / block.y);
19
    matrix_multiplication_v0<<<grid, block>>>(A, B, C, M, N, K);
20
}

展开收起

V1: Original `16 x 16` Shared-Memory Tiling#

Approach#

这个原始版本仍让每个 thread 负责一个输出元素，但把归约维度划分为长度为 16 的 tile。每个 16 x 16 block 负责生成 C 中的一个 16 x 16 tile，从而形成一条明确的数据流动路径：

1
global memory -> shared memory -> thread registers -> C

global memory 保存完整矩阵，shared memory 充当 block 内部的暂存区，用于保存当前参与计算的一对输入 tile；每个 thread 则在 register 中维护自己的标量累加器。

处理 reduction tile tile 时，thread (ty, tx) 负责加载：

1
A[row, tile * 16 + tx] -> A_shared[ty][tx]
2
B[tile * 16 + ty, col] -> B_shared[ty][tx]

全部 256 个 thread 完成加载后，每个 A 值都可由负责不同输出列的 16 个 thread 共享，每个 B 值也可由负责不同输出行的 16 个 thread 共享。随后，每个 thread 执行 16 次乘加运算：

1
S += A_shared[ty][inner] * B_shared[inner][tx]

这样一来，每个从 global memory 读取一次的输入值，在同一个 block 内最多可以参与 16 次乘加运算。对于一个完整的 reduction tile，block 分别从两个输入矩阵加载 16×1616 \times 1616×16 个值，随后利用这些数据完成 16×16×1616 \times 16 \times 1616×16×16 次乘加运算。其算术计算与 V0 完全相同，但大量重复的 global-memory 流量被延迟更低的 shared-memory 访问取代。

当前 Markdown 查看器不支持嵌入式视频。请直接打开[共享内存分块动画](/posts/assets/02_Matrix_Multiplication/video_02.mp4)。

每个 reduction tile 都需要两次同步屏障（barrier）。第一次 barrier 保证协作式加载（cooperative load）全部完成后才开始计算；第二次 barrier 则保证所有 thread 都用完当前 tile 后，下一个 tile 才能覆盖 shared memory 中的现有数据。

最后一个 reduction tile 可能不足 16 个元素。遇到无效的输入坐标时，程序不会提前返回，而是向 shared memory 写入零。零填充（zero padding）不会改变点积结果，同时能让所有 thread 都到达相同的 barrier，避免在同步位置发生分支发散。最终写回输出时，再通过 row < M && col < K 单独进行边界保护。

Solution#

1
#include <cuda_runtime.h>
2

3
#define BLOCK_SIZE  16
4
__global__ void matrix_multiplication_kernel(const float* A, const float* B, float* C, int M, int N,
5
                                             int K) {
6
    __shared__ float A_shared[BLOCK_SIZE][BLOCK_SIZE];
7
    __shared__ float B_shared[BLOCK_SIZE][BLOCK_SIZE];
8
    int row = threadIdx.y + blockIdx.y * blockDim.y;
9
    int col = threadIdx.x + blockIdx.x * blockDim.x;
10
    int tx = threadIdx.x;
11
    int ty = threadIdx.y;
12
    float S = 0.0f;
13
    for (int i = 0 ; i < (N + BLOCK_SIZE - 1) / BLOCK_SIZE ; i++){
14
        int col2 = tx + i * BLOCK_SIZE;
15
        int row2 = ty + i * BLOCK_SIZE;
16
        A_shared[ty][tx] = (row < M && col2 < N) ? A[row * N + col2] : 0.0f;
17
        B_shared[ty][tx] = (row2 < N && col < K) ? B[row2 * K + col] : 0.0f;
18
        __syncthreads();
19
        for (int j = 0 ; j <  BLOCK_SIZE ; j++){
20
            S += A_shared[ty][j] * B_shared[j][tx];
21
        }
22
        __syncthreads();
23
    }
24
    if(row < M && col < K){
25
        C[row * K + col] = S ;
26
    }
27
}
28

29
extern "C" void solve(const float* A, const float* B, float* C, int M, int N, int K) {
30
    dim3 threadsPerBlock(16, 16);
31
    dim3 blocksPerGrid((K + threadsPerBlock.x - 1) / threadsPerBlock.x,
32
                       (M + threadsPerBlock.y - 1) / threadsPerBlock.y);
33

34
    matrix_multiplication_kernel<<<blocksPerGrid, threadsPerBlock>>>(A, B, C, M, N, K);
35
}

展开收起

V1: One-Dimensional Thread Tiling#

Approach#

一维线程分块（1D Thread Tiling）版本使用以下参数：

1
BM = 64
2
BN = 64
3
BK = 16
4
TM = 16

每个 block 生成一个 64 x 64 输出 tile。block 的维度为 (64, 4)，共包含 256 个 thread；每个 thread 负责同一输出列中的 16 行：

1
row = blockIdx.y * 64 + ty * 16 + local_row
2
col = blockIdx.x * 64 + tx

对应的 16 个部分和保存在 float sum[16] 中。在一次归约步骤（reduction step）内，thread 从 shared memory 读取一个 B 值，并在全部 16 个累加器之间复用它：

sum[r]+=SA[ty×TM+r,inner]×b,0≤r<16.\text{sum}[r] \mathrel{+}= SA[ty \times TM+r,\text{inner}] \times b, \quad 0 \le r < 16.sum[r]+=SA[ty×TM+r,inner]×b,0≤r<16.

在每个 thread 只负责一个输出的 kernel 中，每到一个归约位置（reduction position），thread 都要从 shared memory 各读取一个 A 值和一个 B 值，才能执行一次乘加运算。这里，同一个保存 B 值的 register 可以连续支持 16 次乘加。虽然额外的累加器会消耗更多 register，但它们增加了每次 shared-memory 访问所对应的算术工作量，也让 block 级同步的成本分摊到每个 thread 的更多计算上。

当前 Markdown 查看器不支持嵌入式视频。请直接打开[一维线程分块动画](/posts/assets/02_Matrix_Multiplication/video_03.mp4)。

block 仍以协作方式填充 SA[64][16] 和 SB[16][64]。将 thread index 展平后，输入加载便不再依赖 thread 与输出位置之间的归属映射。协作式加载和最终写回都会进行边界检查，因此位于矩阵边缘的不完整 tile 也能得到正确结果。

Solution#

1
#include <cuda_runtime.h>
2

3
template <int BM = 64, int BN = 64, int BK = 16, int TM = 16>
4
__global__ void matrix_multiplication_v1_1d(const float* __restrict__ A,
5
                                            const float* __restrict__ B,
6
                                            float* __restrict__ C,
7
                                            int M, int N, int K) {
8
    constexpr int NT = (BM / TM) * BN;
9
    __shared__ float SA[BM][BK];
10
    __shared__ float SB[BK][BN];
11
    const int tx = threadIdx.x;
12
    const int ty = threadIdx.y;
13
    const int tid = ty * blockDim.x + tx;
14
    const int col = blockIdx.x * BN + tx;
15
    float sum[TM] = {0.0f};
16

17
    for (int tile = 0; tile < N; tile += BK) {
18
        #pragma unroll
19
        for (int i = tid; i < BM * BK; i += NT) {
20
            const int local_row = i / BK;
21
            const int local_col = i % BK;
22
            const int row = blockIdx.y * BM + local_row;
23
            const int inner = tile + local_col;
24
            SA[local_row][local_col] =
25
                (row < M && inner < N) ? A[row * N + inner] : 0.0f;
26
        }
27
        #pragma unroll
28
        for (int i = tid; i < BK * BN; i += NT) {
29
            const int local_row = i / BN;
30
            const int local_col = i % BN;
31
            const int inner = tile + local_row;
32
            const int output_col = blockIdx.x * BN + local_col;
33
            SB[local_row][local_col] =
34
                (inner < N && output_col < K) ? B[inner * K + output_col] : 0.0f;
35
        }
36
        __syncthreads();
37

38
        #pragma unroll
39
        for (int inner = 0; inner < BK; ++inner) {
40
            const float b = SB[inner][tx];
41
            #pragma unroll
42
            for (int row = 0; row < TM; ++row) {
43
                sum[row] += SA[ty * TM + row][inner] * b;
44
            }
45
        }
46
        __syncthreads();
47
    }
48

49
    #pragma unroll
50
    for (int local_row = 0; local_row < TM; ++local_row) {
51
        const int row = blockIdx.y * BM + ty * TM + local_row;
52
        if (row < M && col < K) C[row * K + col] = sum[local_row];
53
    }
54
}
55

56
extern "C" void solve(const float* A, const float* B, float* C, int M, int N, int K) {
57
    constexpr int BM = 64, BN = 64, BK = 16, TM = 16;
58
    const dim3 block(BN, BM / TM);
59
    const dim3 grid((K + BN - 1) / BN, (M + BM - 1) / BM);
60
    matrix_multiplication_v1_1d<BM, BN, BK, TM><<<grid, block>>>(A, B, C, M, N, K);
61
}

展开收起

V1: Two-Dimensional Thread Tiling#

Approach#

二维线程分块（2D Thread Tiling）版本使用以下参数：

1
BM = 64
2
BN = 64
3
BK = 16
4
TM = 8
5
TN = 4

block 的维度为 (16, 8)，即总共 128 个 thread。每个 thread 负责一个 8 x 4 微分块（micro-tile），其中包含 32 个输出元素：

1
row = blockIdx.y * 64 + ty * 8 + local_row
2
col = blockIdx.x * 64 + tx * 4 + local_col

在一个 reduction position 上，thread 会把 8 个 A 值和 4 个 B 值从 shared memory 读入 register，并通过二者的外积（outer product）更新全部 32 个累加器：

Cmicro+=[a0a1⋮a7][b0b1b2b3].C_{\text{micro}} \mathrel{+}= \begin{bmatrix} a_0\\a_1\\\vdots\\a_7 \end{bmatrix} \begin{bmatrix} b_0&b_1&b_2&b_3 \end{bmatrix}.Cmicro+=a0a1⋮a7[b0b1b2b3].

这提供了观察矩阵乘法的另一种视角：单个标量输出可以自然地描述为内积，而整个输出 tile 则可以看作沿 reduction 维度不断累加一系列秩一外积。在某个 reduction position 上，由 8 个元素组成的 A 列片段与由 4 个元素组成的 B 行片段，会共同更新该 thread 负责的整个 8 x 4 输出 tile。

矩阵乘法的外积视角

每个 reduction position 都会为输出 tile 贡献一个秩一外积。

当前 Markdown 查看器不支持嵌入式视频。请直接打开[二维线程分块动画](/posts/assets/02_Matrix_Multiplication/video_04.mp4)。

每个保存 A 值的 register 都会参与 4 个输出列的计算，每个保存 B 值的 register 则会参与 8 个输出行的计算。因此，thread 只需从 shared memory 读取 12 个标量，就可以执行 32 次乘加运算。

这正是二维 thread tiling 的核心优势：必要的数学工作量没有改变，但两个操作数都能先在 register 中得到充分复用，之后才需要再次访问 shared memory。

Solution#

1
#include <cuda_runtime.h>
2

3
template <int BM = 64, int BN = 64, int BK = 16, int TM = 8, int TN = 4>
4
__global__ void matrix_multiplication_v1_2d(const float* __restrict__ A,
5
                                            const float* __restrict__ B,
6
                                            float* __restrict__ C,
7
                                            int M, int N, int K) {
8
    constexpr int NT = (BM / TM) * (BN / TN);
9
    __shared__ float SA[BM][BK];
10
    __shared__ float SB[BK][BN];
11
    const int tx = threadIdx.x;
12
    const int ty = threadIdx.y;
13
    const int tid = ty * blockDim.x + tx;
14
    float sum[TM][TN] = {{0.0f}};
15

16
    for (int tile = 0; tile < (N + BK - 1) / BK; ++tile) {
17
        const int base_inner = tile * BK;
18
        #pragma unroll
19
        for (int i = tid; i < BM * BK; i += NT) {
20
            const int local_row = i / BK;
21
            const int local_col = i % BK;
22
            const int row = blockIdx.y * BM + local_row;
23
            const int inner = base_inner + local_col;
24
            SA[local_row][local_col] =
25
                (row < M && inner < N) ? A[row * N + inner] : 0.0f;
26
        }
27
        #pragma unroll
28
        for (int i = tid; i < BK * BN; i += NT) {
29
            const int local_row = i / BN;
30
            const int local_col = i % BN;
31
            const int inner = base_inner + local_row;
32
            const int col = blockIdx.x * BN + local_col;
33
            SB[local_row][local_col] =
34
                (inner < N && col < K) ? B[inner * K + col] : 0.0f;
35
        }
36
        __syncthreads();
37

38
        #pragma unroll
39
        for (int inner = 0; inner < BK; ++inner) {
40
            float a_reg[TM], b_reg[TN];
41
            #pragma unroll
42
            for (int row = 0; row < TM; ++row) a_reg[row] = SA[ty * TM + row][inner];
43
            #pragma unroll
44
            for (int col = 0; col < TN; ++col) b_reg[col] = SB[inner][tx * TN + col];
45
            #pragma unroll
46
            for (int row = 0; row < TM; ++row)
47
                #pragma unroll
48
                for (int col = 0; col < TN; ++col)
49
                    sum[row][col] += a_reg[row] * b_reg[col];
50
        }
51
        __syncthreads();
52
    }
53

54
    #pragma unroll
55
    for (int local_row = 0; local_row < TM; ++local_row) {
56
        const int row = blockIdx.y * BM + ty * TM + local_row;
57
        if (row >= M) continue;
58
        #pragma unroll
59
        for (int local_col = 0; local_col < TN; ++local_col) {
60
            const int col = blockIdx.x * BN + tx * TN + local_col;
61
            if (col < K) C[row * K + col] = sum[local_row][local_col];
62
        }
63
    }
64
}
65

66
extern "C" void solve(const float* A, const float* B, float* C, int M, int N, int K) {
67
    constexpr int BM = 64, BN = 64, BK = 16, TM = 8, TN = 4;
68
    const dim3 block(BN / TN, BM / TM);
69
    const dim3 grid((K + BN - 1) / BN, (M + BM - 1) / BM);
70
    matrix_multiplication_v1_2d<BM, BN, BK, TM, TN><<<grid, block>>>(A, B, C, M, N, K);
71
}

展开收起

V2: Vectorized `float4` Transfers (`BK = 32`)#

Approach#

V2 保留了 8 x 4 register micro-tile，并把 reduction tile 的深度改为 BK = 32。一个 float4 包含 4 个连续的 FP32 值，共占 16 字节。kernel 在以下数据传输位置使用它，一次处理 4 个相邻值：

从 global memory 加载 A 和 B；
将数据写入 shared memory；
从 shared memory 读取相邻的 B 值；
最终写入 C 的相邻列。

只有当访问地址满足对齐要求，并且对应的 4 个元素全部有效时，向量加载才是合法的。因此，仅当 N 和 K 都能被 4 整除时，程序才会选择向量化 kernel；任一条件不满足，solve(...) 都会启动能够安全处理边界的标量实现。通用的 LeetGPU 接口必须保留这一回退路径（fallback），不能为了向量化而牺牲不规则矩阵形状（shape）的正确性。

在向量化路径中，最后一个 block 的输出行仍需单独进行边界保护。由于此时 K 是 4 的倍数，一个由 4 列组成的分组要么完整落在矩阵范围内，要么完整位于矩阵范围外，不会只剩部分列有效。

行主序布局让同一行中的连续列在内存中相邻，因此适合沿 A 的归约坐标、B 的输出列坐标以及 C 的输出列坐标进行向量化。在这些传输位置，一条向量指令可以替代 4 条标量指令。实际传输的总字节数，以及矩阵乘法所需的 2MNK2MNK2MNK 次浮点运算都没有改变；向量化主要减少的是指令数量和地址生成开销。因此，实际收益取决于地址对齐、指令发射效率，以及数据移动与算术计算之间的平衡，并不会仅仅因为使用了更宽的数据类型就自动出现。

Solution#

1
#include <cuda_runtime.h>
2

3
#ifndef MATMUL_V2_BM
4
#define MATMUL_V2_BM 64
5
#endif
6
#ifndef MATMUL_V2_BN
7
#define MATMUL_V2_BN 64
8
#endif
9
#ifndef MATMUL_V2_BK
10
#define MATMUL_V2_BK 32
11
#endif
12
#ifndef MATMUL_V2_TM
13
#define MATMUL_V2_TM 8
14
#endif
15
#ifndef MATMUL_V2_TN
16
#define MATMUL_V2_TN 4
17
#endif
18

19
__global__ void scalar_fallback(const float* A, const float* B, float* C,
20
                                int M, int N, int K) {
21
    const int row = blockIdx.y * blockDim.y + threadIdx.y;
22
    const int col = blockIdx.x * blockDim.x + threadIdx.x;
23
    if (row >= M || col >= K) return;
24
    float sum = 0.0f;
25
    for (int inner = 0; inner < N; ++inner) {
26
        sum += A[row * N + inner] * B[inner * K + col];
27
    }
28
    C[row * K + col] = sum;
29
}
30

31
template <int BM = MATMUL_V2_BM, int BN = MATMUL_V2_BN,
32
          int BK = MATMUL_V2_BK, int TM = MATMUL_V2_TM,
33
          int TN = MATMUL_V2_TN>
34
__global__ void matrix_multiplication_v2(const float* __restrict__ A,
35
                                         const float* __restrict__ B,
36
                                         float* __restrict__ C,
37
                                         int M, int N, int K) {
38
    static_assert(BK % 4 == 0 && BN % 4 == 0 && TN % 4 == 0,
39
                  "Vectorized dimensions must be multiples of four");
40
    constexpr int NT = (BM / TM) * (BN / TN);
41
    __shared__ float SA[BM][BK];
42
    __shared__ float SB[BK][BN];
43
    const int tx = threadIdx.x;
44
    const int ty = threadIdx.y;
45
    const int tid = ty * blockDim.x + tx;
46
    float sum[TM][TN] = {{0.0f}};
47

48
    for (int tile = 0; tile < (N + BK - 1) / BK; ++tile) {
49
        const int base_inner = tile * BK;
50
        #pragma unroll
51
        for (int i = tid; i < BM * BK / 4; i += NT) {
52
            const int local_row = i / (BK / 4);
53
            const int local_col = (i % (BK / 4)) * 4;
54
            const int row = blockIdx.y * BM + local_row;
55
            const int inner = base_inner + local_col;
56
            float4 value = make_float4(0, 0, 0, 0);
57
            if (row < M && inner + 3 < N) {
58
                value = *reinterpret_cast<const float4*>(&A[row * N + inner]);
59
            }
60
            *reinterpret_cast<float4*>(&SA[local_row][local_col]) = value;
61
        }
62
        #pragma unroll
63
        for (int i = tid; i < BK * BN / 4; i += NT) {
64
            const int local_row = i / (BN / 4);
65
            const int local_col = (i % (BN / 4)) * 4;
66
            const int inner = base_inner + local_row;
67
            const int col = blockIdx.x * BN + local_col;
68
            float4 value = make_float4(0, 0, 0, 0);
69
            if (inner < N && col + 3 < K) {
70
                value = *reinterpret_cast<const float4*>(&B[inner * K + col]);
71
            }
72
            *reinterpret_cast<float4*>(&SB[local_row][local_col]) = value;
73
        }
74
        __syncthreads();
75

76
        #pragma unroll
77
        for (int inner = 0; inner < BK; ++inner) {
78
            const float4 b = *reinterpret_cast<const float4*>(&SB[inner][tx * TN]);
79
            #pragma unroll
80
            for (int row = 0; row < TM; ++row) {
81
                const float a = SA[ty * TM + row][inner];
82
                sum[row][0] += a * b.x;
83
                sum[row][1] += a * b.y;
84
                sum[row][2] += a * b.z;
85
                sum[row][3] += a * b.w;
86
            }
87
        }
88
        __syncthreads();
89
    }
90

91
    #pragma unroll
92
    for (int local_row = 0; local_row < TM; ++local_row) {
93
        const int row = blockIdx.y * BM + ty * TM + local_row;
94
        const int col = blockIdx.x * BN + tx * TN;
95
        if (row < M && col + 3 < K) {
96
            *reinterpret_cast<float4*>(&C[row * K + col]) =
97
                make_float4(sum[local_row][0], sum[local_row][1],
98
                            sum[local_row][2], sum[local_row][3]);
99
        }
100
    }
101
}
102

103
extern "C" void solve(const float* A, const float* B, float* C, int M, int N, int K) {
104
    if ((N & 3) != 0 || (K & 3) != 0) {
105
        const dim3 block(16, 16);
106
        const dim3 grid((K + 15) / 16, (M + 15) / 16);
107
        scalar_fallback<<<grid, block>>>(A, B, C, M, N, K);
108
    } else {
109
        constexpr int BM = MATMUL_V2_BM, BN = MATMUL_V2_BN;
110
        constexpr int BK = MATMUL_V2_BK, TM = MATMUL_V2_TM;
111
        constexpr int TN = MATMUL_V2_TN;
112
        const dim3 block(BN / TN, BM / TM);
113
        const dim3 grid((K + BN - 1) / BN, (M + BM - 1) / BM);
114
        matrix_multiplication_v2<BM, BN, BK, TM, TN>
115
            <<<grid, block>>>(A, B, C, M, N, K);
116
    }
117
}

展开收起

V2 Control: Vectorized Single Buffer (`BK = 16`)#

Approach#

这个对照版本仍编译同一个单缓冲（single-buffered）V2 kernel，但将 BK = 16，以匹配 V3 的归约分块深度（reduction-tile depth）。V2 的其余设计均保持不变，包括每个 thread 负责的输出、向量化传输、register micro-tile、边界 fallback 和启动布局（launch geometry）。

这一参数匹配的配置，可以把 reduction-tile depth 的影响与双缓冲（double buffering）的影响区分开来。在两款受测 GPU 上，single-buffered BK = 16 对照版本仍然快于 V3。因此，不能仅凭 V2 原本使用 BK = 32、而 V3 使用 BK = 16 来解释 V3 较低的实测性能；在当前实现中，额外的预取 register、两套 shared-memory buffer、buffer 切换与同步操作，同样没有带来足以抵消自身成本的收益。

Solution#

1
#define MATMUL_V2_BK 16
2

3
#include "02_Matrix_Multiplication_v2_vectorized.cu"

V2 Control: Vectorized Single Buffer (`128 x 128`)#

Approach#

这个对照版本保留 V2 的 single-buffered 算法，只把编译期 tile 参数改为 BM = 128、BN = 128 和 BK = 16。每个 thread 负责的 8 x 4 micro-tile、float4 传输、标量边界 fallback 以及累加逻辑均保持不变。

block 的规模从 128 个 thread 增加到 512 个 thread，生成的输出元素数量也变为原来的 4 倍；然而，更大的 tile 并未提升这个 single-buffered kernel 的性能。在两款受测 GPU 上，它都慢于参数匹配的 64 x 64、BK = 16 对照版本。这说明额外的 block 级数据复用本身还不足以抵消更大 block 带来的调度成本和资源开销。

Solution#

1
#define MATMUL_V2_BM 128
2
#define MATMUL_V2_BN 128
3
#define MATMUL_V2_BK 16
4
#define MATMUL_V2_TM 8
5
#define MATMUL_V2_TN 4
6

7
#include "02_Matrix_Multiplication_v2_vectorized.cu"

V3: Register Prefetch and Double-Buffered Shared Memory#

Base `64 x 64` Configuration#

Approach#

V3 使用以下参数：

1
BM = 64
2
BN = 64
3
BK = 16
4
TM = 8
5
TN = 4

它为每个 shared-memory tile 分配两个副本：

1
SA[2][64][16]
2
SB[2][16][64]

第一个 tile 在主循环之前加载。这一阶段称为流水线序言（pipeline prologue）：只有一对完整的输入 tile 已进入 shared memory 后，计算才能开始。

对于后续的每个 tile，thread 首先发出 global load，将数据加载到私有的预取 register。随后，它基于当前 shared-memory buffer 进行计算，等待所有 thread 都用完该 buffer，再将预取值写入另一个 buffer，最后在交换读写 buffer index 之前完成同步。这一阶段称为稳态（steady state），其中 reduction tile iii 的计算与 tile i+1i+1i+1 的数据移动交错进行。

最后一个预取 tile 成为当前 tile 后，kernel 会完成其计算，但不再请求下一个 tile。这一阶段称为流水线尾声（pipeline epilogue）。

预期的流水线如下：

1
tile i + 1: global memory -> prefetch registers
2
tile i:     shared memory -> accumulator registers
3
tile i + 1: prefetch registers -> alternate shared-memory buffer

这是带 register prefetch 的软件双缓冲（software double buffering），而不是异步 cp.async 流水线。global load 会在当前 tile 被使用之前发出，从而使编译器和 GPU 有机会将尚未完成的内存操作与独立的算术计算重叠执行。不过，当 shared buffer 切换角色时，该实现仍需要 block 范围的 barrier。

实际源码使用直接的行主序 tile 布局保存 SA[buffer][row][inner] 和 SB[buffer][inner][col]，并未暗中加入转置 shared-memory 布局或其他未实现的访问变换。这一点很重要，因为评价一条流水线时，不能脱离其实际加载模式、register footprint、shared-memory layout 和同步调度。

double buffering 并不会自动提升速度。它使 shared-memory 存储量翻倍，并增加预取 register、控制流、pipeline prologue 和 epilogue。这些开销可能降低 occupancy，或抵消隐藏延迟带来的收益。在实测工作负载中，基础 64 x 64 V3 在两款受测 GPU 上都略慢于参数匹配的 single-buffer V2 对照版本。后面的 128 x 128 V3 配置更快，是因为更大的输出 tile 同时改变了复用和调度；保留各个独立对照版本，可以避免将该结果错误地完全归因于 double buffering。

Solution#

1
#include <cuda_runtime.h>
2

3
#ifndef MATMUL_BM
4
#define MATMUL_BM 64
5
#define MATMUL_BN 64
6
#define MATMUL_BK 16
7
#define MATMUL_TM 8
8
#define MATMUL_TN 4
9
#endif
10

11
__global__ void scalar_fallback_v3(const float* A, const float* B, float* C,
12
                                   int M, int N, int K) {
13
    const int row = blockIdx.y * blockDim.y + threadIdx.y;
14
    const int col = blockIdx.x * blockDim.x + threadIdx.x;
15
    if (row >= M || col >= K) return;
16
    float sum = 0.0f;
17
    for (int inner = 0; inner < N; ++inner) {
18
        sum += A[row * N + inner] * B[inner * K + col];
19
    }
20
    C[row * K + col] = sum;
21
}
22

23
template <int BM = 64, int BN = 64, int BK = 16, int TM = 8, int TN = 4>
24
__global__ void matrix_multiplication_v3(const float* __restrict__ A,
25
                                         const float* __restrict__ B,
26
                                         float* __restrict__ C,
27
                                         int M, int N, int K) {
28
    constexpr int NT = (BM / TM) * (BN / TN);
29
    constexpr int A_FLOAT4_PER_THREAD = (BM * BK / 4 + NT - 1) / NT;
30
    constexpr int B_FLOAT4_PER_THREAD = (BK * BN / 4 + NT - 1) / NT;
31
    __shared__ float SA[2][BM][BK];
32
    __shared__ float SB[2][BK][BN];
33

34
    const int tx = threadIdx.x;
35
    const int ty = threadIdx.y;
36
    const int tid = ty * blockDim.x + tx;
37
    float a_prefetch[A_FLOAT4_PER_THREAD][4];
38
    float b_prefetch[B_FLOAT4_PER_THREAD][4];
39
    float sum[TM][TN] = {{0.0f}};
40

41
    auto load_tile = [&](int tile) {
42
        const int base_inner = tile * BK;
43
        #pragma unroll
44
        for (int i = 0; i < A_FLOAT4_PER_THREAD; ++i) {
45
            const int index = (i * NT + tid) * 4;
46
            float4 value = make_float4(0, 0, 0, 0);
47
            if (index < BM * BK) {
48
                const int local_row = index / BK;
49
                const int local_col = index % BK;
50
                const int row = blockIdx.y * BM + local_row;
51
                const int inner = base_inner + local_col;
52
                if (row < M && inner + 3 < N) {
53
                    value = *reinterpret_cast<const float4*>(&A[row * N + inner]);
54
                }
55
            }
56
            a_prefetch[i][0] = value.x; a_prefetch[i][1] = value.y;
57
            a_prefetch[i][2] = value.z; a_prefetch[i][3] = value.w;
58
        }
59
        #pragma unroll
60
        for (int i = 0; i < B_FLOAT4_PER_THREAD; ++i) {
61
            const int index = (i * NT + tid) * 4;
62
            float4 value = make_float4(0, 0, 0, 0);
63
            if (index < BK * BN) {
64
                const int local_row = index / BN;
65
                const int local_col = index % BN;
66
                const int inner = base_inner + local_row;
67
                const int col = blockIdx.x * BN + local_col;
68
                if (inner < N && col + 3 < K) {
69
                    value = *reinterpret_cast<const float4*>(&B[inner * K + col]);
70
                }
71
            }
72
            b_prefetch[i][0] = value.x; b_prefetch[i][1] = value.y;
73
            b_prefetch[i][2] = value.z; b_prefetch[i][3] = value.w;
74
        }
75
    };
76

77
    auto store_tile = [&](int buffer) {
78
        #pragma unroll
79
        for (int i = 0; i < A_FLOAT4_PER_THREAD; ++i) {
80
            const int index = (i * NT + tid) * 4;
81
            if (index < BM * BK) {
82
                const int local_row = index / BK;
83
                const int local_col = index % BK;
84
                *reinterpret_cast<float4*>(&SA[buffer][local_row][local_col]) =
85
                    make_float4(a_prefetch[i][0], a_prefetch[i][1],
86
                                a_prefetch[i][2], a_prefetch[i][3]);
87
            }
88
        }
89
        #pragma unroll
90
        for (int i = 0; i < B_FLOAT4_PER_THREAD; ++i) {
91
            const int index = (i * NT + tid) * 4;
92
            if (index < BK * BN) {
93
                const int local_row = index / BN;
94
                const int local_col = index % BN;
95
                *reinterpret_cast<float4*>(&SB[buffer][local_row][local_col]) =
96
                    make_float4(b_prefetch[i][0], b_prefetch[i][1],
97
                                b_prefetch[i][2], b_prefetch[i][3]);
98
            }
99
        }
100
    };
101

102
    const int tiles = (N + BK - 1) / BK;
103
    load_tile(0);
104
    store_tile(0);
105
    __syncthreads();
106

107
    int read_buffer = 0;
108
    for (int tile = 0; tile < tiles; ++tile) {
109
        if (tile + 1 < tiles) load_tile(tile + 1);
110

111
        #pragma unroll
112
        for (int inner = 0; inner < BK; ++inner) {
113
            #pragma unroll
114
            for (int local_col = 0; local_col < TN; local_col += 4) {
115
                const float4 b = *reinterpret_cast<const float4*>(
116
                    &SB[read_buffer][inner][tx * TN + local_col]);
117
                #pragma unroll
118
                for (int row = 0; row < TM; ++row) {
119
                    const float a = SA[read_buffer][ty * TM + row][inner];
120
                    sum[row][local_col + 0] += a * b.x;
121
                    sum[row][local_col + 1] += a * b.y;
122
                    sum[row][local_col + 2] += a * b.z;
123
                    sum[row][local_col + 3] += a * b.w;
124
                }
125
            }
126
        }
127
        __syncthreads();
128

129
        if (tile + 1 < tiles) {
130
            const int write_buffer = read_buffer ^ 1;
131
            store_tile(write_buffer);
132
            __syncthreads();
133
            read_buffer = write_buffer;
134
        }
135
    }
136

137
    #pragma unroll
138
    for (int local_row = 0; local_row < TM; ++local_row) {
139
        const int row = blockIdx.y * BM + ty * TM + local_row;
140
        #pragma unroll
141
        for (int local_col = 0; local_col < TN; local_col += 4) {
142
            const int col = blockIdx.x * BN + tx * TN + local_col;
143
            if (row < M && col + 3 < K) {
144
                *reinterpret_cast<float4*>(&C[row * K + col]) =
145
                    make_float4(sum[local_row][local_col + 0],
146
                                sum[local_row][local_col + 1],
147
                                sum[local_row][local_col + 2],
148
                                sum[local_row][local_col + 3]);
149
            }
150
        }
151
    }
152
}
153

154
#ifndef MATMUL_NO_SOLVE
155
extern "C" void solve(const float* A, const float* B, float* C, int M, int N, int K) {
156
    if ((N & 3) != 0 || (K & 3) != 0) {
157
        const dim3 block(16, 16);
158
        const dim3 grid((K + 15) / 16, (M + 15) / 16);
159
        scalar_fallback_v3<<<grid, block>>>(A, B, C, M, N, K);
160
    } else {
161
        constexpr int BM = MATMUL_BM, BN = MATMUL_BN, BK = MATMUL_BK;
162
        constexpr int TM = MATMUL_TM, TN = MATMUL_TN;
163
        const dim3 block(BN / TN, BM / TM);
164
        const dim3 grid((K + BN - 1) / BN, (M + BM - 1) / BM);
165
        matrix_multiplication_v3<BM, BN, BK, TM, TN>
166
            <<<grid, block>>>(A, B, C, M, N, K);
167
    }
168
}
169
#endif

展开收起

Large-Tile `128 x 128` Configuration#

Approach#

large-tile 配置将 V3 的编译期参数改为：

1
BM = 128
2
BN = 128
3
BK = 16
4
TM = 8
5
TN = 4

block 包含 (128 / 4) x (128 / 8) = 32 x 16 = 512 个 thread。每个 thread 仍负责 32 个输出元素，而整个 block 共同生成一个 128 x 128 tile。其 double-buffered 输入存储占用 32 KiB：

1
2 * (128 * 16 + 16 * 128) * sizeof(float)

更大的输出 tile 提高了 block 级复用，并将 input-tile 加载开销分摊到更多输出元素上，同时完整保留基础 V3 配置的 register-prefetch 和 double-buffer 机制。它也增加了 block 的 thread 数量，并改变 register 需求、occupancy 和调度，因此它是 V3 的一种参数配置，而不是独立的优化原理。

在修订后的 benchmark 中，该配置是两款受测 GPU 上本项目最快的 CUDA 实现。然而，参数匹配的 single-buffered 实验在扩大到相同的 128 x 128 tile 后反而变慢。因此，收益并非单独来自 tile 大小：只有当更大的输出 tile 与这条 double-buffered 数据路径结合时，收益才会出现。实测结果表明这两种选择存在相互作用，而不能由此得出更大 tile 普遍更快的结论。

Solution#

1
#define MATMUL_BM 128
2
#define MATMUL_BN 128
3
#define MATMUL_BK 16
4
#define MATMUL_TM 8
5
#define MATMUL_TN 4
6

7
#include "02_Matrix_Multiplication_v3_double_buffered.cu"

Controlled Comparison#

下表 4 行均使用 BK = 16、每个 thread 一个 8 x 4 micro-tile、确定性的非零输入、5 次 warm-up iteration 和 10 次 measured iteration。

Buffering	Output Tile	RTX 5070 Ti Laptop GPU	RTX 4090
Single buffer	`64 x 64`	38.761 ms	11.314 ms
Single buffer	`128 x 128`	38.842 ms	13.146 ms
Double buffer	`64 x 64`	39.642 ms	11.514 ms
Double buffer	`128 x 128`	34.590 ms	10.052 ms

在 64 x 64 配置下，double buffering 的影响很小且取决于平台：在 RTX 4090 上略快，在 RTX 5070 Ti Laptop GPU 上略慢。将 single-buffered kernel 扩大到 128 x 128 后，两款 GPU 上的速度都更慢。相比之下，128 x 128 double-buffered 配置明显快于 single-buffered large tile，也是表中最快的通用配置。因此，实测加速来自更大输出 tile 与 double-buffered 数据路径的相互作用，不能由此得出任一选择单独使用时普遍更快的结论。

以下编译器资源报告由面向 RTX 4090 target 的 nvcc -O2 -arch=sm_89 -Xptxas=-v 生成：

Buffering	Output Tile	Threads per Block	Registers per Thread	Shared Memory per Block
Single buffer	`64 x 64`	128	86	8 KiB
Single buffer	`128 x 128`	512	86	16 KiB
Double buffer	`64 x 64`	128	122	16 KiB
Double buffer	`128 x 128`	512	100	32 KiB

该报告证实了资源方面的权衡。single-buffered large tile 仍为每个 thread 使用 86 个 register，但 thread 数量增加到 4 倍，因此每个 block 分配的 register 从 11,008 个增加到 44,032 个。基础 double-buffered kernel 因预取状态而将每个 thread 的 register 用量提高到 122 个。对于 large double-buffered 实例，编译器为每个 thread 使用 100 个 register，且没有 spill；block 在保持相同 per-thread micro-tile 的同时，完成了 4 倍的输出工作。这些事实解释了为何 occupancy 和复用程度会随配置发生显著变化，但仅凭 register 与 shared-memory 数量并不能确定实际运行时的 stall 行为。

V4: Hardware-Asynchronous `cp.async` Pipeline#

Approach#

V4 保留 large V3 配置的 128 x 128 输出 tile、BK = 16 和每个 thread 一个 8 x 4 micro-tile，但改变了下一个输入 tile 进入 shared memory 的方式。

V3 首先将 global-memory 中的值加载到普通 thread register，随后再把这些 register 中的值存入另一个 shared-memory buffer。V4 则改为发出 Ampere 或更新架构支持的 PTX 指令：

1
cp.async.cg.shared.global

每条指令直接从 global memory 向 shared memory 复制 16 字节。cp.async.commit_group 提交一个 thread 发出的复制操作，cp.async.wait_group 0 则等待 block 所需的全部未完成 group 执行完毕。

该流水线包含 prologue、steady state 和 epilogue：

prologue 异步填充 shared buffer 0，并在首次使用前等待填充完成。
在一次 steady-state iteration 中，block 会先将 tile i + 1 的复制操作发往另一个 buffer，再计算 tile i。
算术循环只读取当前 buffer，因此 GPU 可将 fused multiply-add 工作与尚未完成的 global-to-shared 复制重叠执行。
buffer 交换角色前，cp.async.wait_group 0 保证复制完成，__syncthreads() 则使新 tile 对所有 thread 可见。
最后一个 tile 没有后继，因此循环结束时不会再发出另一组复制操作。

越界的 16 字节 vector 会使用有效字节数为零的 source-size operand。硬件将目标 shared-memory vector 填零，而不会解引用无效的 global address。reduction 维度或 output-column 维度不适合 4 个 float 向量传输的输入会使用标量 fallback。compute capability 低于 8.0 的设备因不支持 cp.async，也会使用 fallback。

与 V3 的 register-prefetch 软件流水线不同，这是一条真正的硬件异步流水线（hardware-asynchronous pipeline）。不过，它仍不保证更快：两个版本在 shared buffer 切换角色时都需要同步，而包含 512 个 thread 的 block、地址生成、copy-group 管理和算术指令流都会争用执行资源。

Solution#

1
#include <cuda_runtime.h>
2

3
namespace {
4

5
constexpr int BM = 128;
6
constexpr int BN = 128;
7
constexpr int BK = 16;
8
constexpr int TM = 8;
9
constexpr int TN = 4;
10
constexpr int THREADS_X = BN / TN;
11
constexpr int THREADS_Y = BM / TM;
12
constexpr int THREAD_COUNT = THREADS_X * THREADS_Y;
13

14
__global__ void scalar_fallback_v4(const float* A, const float* B, float* C,
15
                                   int M, int N, int K) {
16
    const int row = blockIdx.y * blockDim.y + threadIdx.y;
17
    const int col = blockIdx.x * blockDim.x + threadIdx.x;
18
    if (row >= M || col >= K) return;
19

20
    float sum = 0.0f;
21
    for (int inner = 0; inner < N; ++inner) {
22
        sum += A[row * N + inner] * B[inner * K + col];
23
    }
24
    C[row * K + col] = sum;
25
}
26

27
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
28
__device__ __forceinline__ void cp_async_16(void* shared_destination,
29
                                            const void* global_source,
30
                                            int valid_bytes) {
31
    const unsigned int shared_address =
32
        static_cast<unsigned int>(__cvta_generic_to_shared(shared_destination));
33
    asm volatile(
34
        "cp.async.cg.shared.global [%0], [%1], 16, %2;\n"
35
        :
36
        : "r"(shared_address), "l"(global_source), "r"(valid_bytes));
37
}
38

39
__device__ __forceinline__ void cp_async_commit() {
40
    asm volatile("cp.async.commit_group;\n" : :);
41
}
42

43
__device__ __forceinline__ void cp_async_wait_all() {
44
    asm volatile("cp.async.wait_group 0;\n" : :);
45
}
46
#endif
47

48
__global__ void matrix_multiplication_v4_cp_async(
49
    const float* __restrict__ A,
50
    const float* __restrict__ B,
51
    float* __restrict__ C,
52
    int M, int N, int K) {
53
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
54
    __shared__ __align__(16) float shared_a[2][BM][BK];
55
    __shared__ __align__(16) float shared_b[2][BK][BN];
56

57
    const int tx = threadIdx.x;
58
    const int ty = threadIdx.y;
59
    const int tid = ty * blockDim.x + tx;
60
    float accumulators[TM][TN] = {{0.0f}};
61

62
    auto issue_tile = [&](int tile, int buffer) {
63
        constexpr int A_VECTORS = BM * BK / 4;
64
        constexpr int B_VECTORS = BK * BN / 4;
65
        const int inner_base = tile * BK;
66

67
        for (int vector_index = tid;
68
             vector_index < A_VECTORS;
69
             vector_index += THREAD_COUNT) {
70
            const int element_index = vector_index * 4;
71
            const int local_row = element_index / BK;
72
            const int local_inner = element_index % BK;
73
            const int row = blockIdx.y * BM + local_row;
74
            const int inner = inner_base + local_inner;
75
            const int valid_bytes = (row < M && inner + 3 < N) ? 16 : 0;
76
            const float* source = A + static_cast<size_t>(row < M ? row : 0) * N +
77
                                  (inner + 3 < N ? inner : 0);
78
            cp_async_16(&shared_a[buffer][local_row][local_inner],
79
                        source, valid_bytes);
80
        }
81

82
        for (int vector_index = tid;
83
             vector_index < B_VECTORS;
84
             vector_index += THREAD_COUNT) {
85
            const int element_index = vector_index * 4;
86
            const int local_inner = element_index / BN;
87
            const int local_col = element_index % BN;
88
            const int inner = inner_base + local_inner;
89
            const int col = blockIdx.x * BN + local_col;
90
            const int valid_bytes = (inner < N && col + 3 < K) ? 16 : 0;
91
            const float* source = B + static_cast<size_t>(inner < N ? inner : 0) * K +
92
                                  (col + 3 < K ? col : 0);
93
            cp_async_16(&shared_b[buffer][local_inner][local_col],
94
                        source, valid_bytes);
95
        }
96
        cp_async_commit();
97
    };
98

99
    const int tile_count = (N + BK - 1) / BK;
100
    issue_tile(0, 0);
101
    cp_async_wait_all();
102
    __syncthreads();
103

104
    int read_buffer = 0;
105
    for (int tile = 0; tile < tile_count; ++tile) {
106
        const bool has_next_tile = tile + 1 < tile_count;
107
        if (has_next_tile) {
108
            issue_tile(tile + 1, read_buffer ^ 1);
109
        }
110

111
        #pragma unroll
112
        for (int inner = 0; inner < BK; ++inner) {
113
            const float4 b = *reinterpret_cast<const float4*>(
114
                &shared_b[read_buffer][inner][tx * TN]);
115
            #pragma unroll
116
            for (int local_row = 0; local_row < TM; ++local_row) {
117
                const float a =
118
                    shared_a[read_buffer][ty * TM + local_row][inner];
119
                accumulators[local_row][0] += a * b.x;
120
                accumulators[local_row][1] += a * b.y;
121
                accumulators[local_row][2] += a * b.z;
122
                accumulators[local_row][3] += a * b.w;
123
            }
124
        }
125

126
        if (has_next_tile) {
127
            cp_async_wait_all();
128
            __syncthreads();
129
            read_buffer ^= 1;
130
        }
131
    }
132

133
    #pragma unroll
134
    for (int local_row = 0; local_row < TM; ++local_row) {
135
        const int row = blockIdx.y * BM + ty * TM + local_row;
136
        const int col = blockIdx.x * BN + tx * TN;
137
        if (row < M && col + 3 < K) {
138
            *reinterpret_cast<float4*>(&C[static_cast<size_t>(row) * K + col]) =
139
                make_float4(accumulators[local_row][0],
140
                            accumulators[local_row][1],
141
                            accumulators[local_row][2],
142
                            accumulators[local_row][3]);
143
        }
144
    }
145
#endif
146
}
147

148
}  // namespace
149

150
extern "C" void solve(const float* A, const float* B, float* C,
151
                      int M, int N, int K) {
152
    static const bool supports_cp_async = [] {
153
        int device = 0;
154
        cudaDeviceProp properties{};
155
        cudaGetDevice(&device);
156
        cudaGetDeviceProperties(&properties, device);
157
        return properties.major >= 8;
158
    }();
159

160
    if (!supports_cp_async || (N & 3) != 0 || (K & 3) != 0) {
161
        const dim3 block(16, 16);
162
        const dim3 grid((K + 15) / 16, (M + 15) / 16);
163
        scalar_fallback_v4<<<grid, block>>>(A, B, C, M, N, K);
164
        return;
165
    }
166

167
    const dim3 block(THREADS_X, THREADS_Y);
168
    const dim3 grid((K + BN - 1) / BN, (M + BM - 1) / BM);
169
    matrix_multiplication_v4_cp_async<<<grid, block>>>(A, B, C, M, N, K);
170
}

展开收起

V5: WMMA TF32 Tensor Core Kernel#

Approach#

V5 同时改变了执行单元和数值格式。公共接口仍然接收和返回 FP32 矩阵，但 fast path 会将输入值显式舍入为 TensorFloat-32，并在 Tensor Core 上使用 FP32 累加器执行矩阵乘加运算。

TF32 保留 FP32 的 8 位指数范围，但使用 10 位显式尾数。转换通过以下指令完成：

1
cvt.rna.tf32.f32

这意味着 V5 在数值上与严格的 FP32 CUDA Core 版本并不完全相同。它以降低输入尾数精度为代价，换取使用 Tensor Core 矩阵指令的能力。因此，测试采用 5e-3 的缩放容差，同时仍会拒绝所有非有限结果，并将每个输出元素与 CPU reference 进行比较。

一个 block 生成一个 128 x 128 输出 tile，并以 32 为单位分块归约 K。该 block 包含 16 个 warp。每个 warp 负责一条 16 x 64 输出带，由 4 个 16 x 16 accumulator fragment 表示。两个 warp 分别覆盖同一个 16 行输出带的左半部分和右半部分。

对于每个 K=32 shared-memory tile：

全部 512 个 thread 协作加载 A[128, 32] 和 B[32, 128]，并将每个值转换为 TF32。
一次 block 同步使两个 tile 均可见。
由于 TF32 WMMA 使用 16 x 16 x 8 fragment，该 tile 被划分为 4 个 WMMA reduction step。
每个 warp 在每个 step 加载一个 A fragment 和 4 个 B fragment。
wmma::mma_sync 更新 4 个 FP32 accumulator fragment。
所有 reduction tile 完成后，wmma::store_matrix_sync 将 fragment 写入行主序 FP32 output memory。

fast path 要求 M 和 K 与 128 x 128 输出 tile 对齐，N 与 BK=32 对齐。其他 shape、较旧架构和边界情况会使用标量 FP32 fallback。这样既保留了题目接口和一般情况下的正确性，也使 WMMA kernel 无需处理 partial-fragment store。

Tensor Core 能够提供很高的峰值吞吐率，但该实现也需要承担显式 TF32 转换、shared-memory staging、频繁 fragment load、block 同步和 WMMA fragment 归属管理等开销。因此，这些测量结果证明的是一个正确的底层 WMMA 实现，并不意味着它必然优于经过精细调优的 CUDA Core kernel。

Solution#

1
#include <cuda_runtime.h>
2
#include <mma.h>
3

4
namespace {
5

6
constexpr int BM = 128;
7
constexpr int BN = 128;
8
constexpr int BK = 32;
9
constexpr int WARPS_PER_BLOCK = 16;
10
constexpr int THREADS_PER_BLOCK = WARPS_PER_BLOCK * 32;
11
constexpr int WMMA_K = 8;
12
constexpr int OUTPUT_FRAGMENTS_PER_WARP = BN / 32;
13

14
__device__ __forceinline__ float to_tf32(float value) {
15
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
16
    unsigned int tf32_bits;
17
    asm("cvt.rna.tf32.f32 %0, %1;" : "=r"(tf32_bits) : "f"(value));
18
    return __uint_as_float(tf32_bits);
19
#else
20
    return value;
21
#endif
22
}
23

24
__global__ void scalar_fallback_v5(const float* A, const float* B, float* C,
25
                                   int M, int N, int K) {
26
    const int row = blockIdx.y * blockDim.y + threadIdx.y;
27
    const int col = blockIdx.x * blockDim.x + threadIdx.x;
28
    if (row >= M || col >= K) return;
29

30
    float sum = 0.0f;
31
    for (int inner = 0; inner < N; ++inner) {
32
        sum += A[row * N + inner] * B[inner * K + col];
33
    }
34
    C[row * K + col] = sum;
35
}
36

37
__global__ void matrix_multiplication_v5_wmma_tf32(
38
    const float* __restrict__ A,
39
    const float* __restrict__ B,
40
    float* __restrict__ C,
41
    int M, int N, int K) {
42
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
43
    using namespace nvcuda;
44

45
    __shared__ __align__(16) float shared_a[BM][BK];
46
    __shared__ __align__(16) float shared_b[BK][BN];
47

48
    const int tid = threadIdx.x;
49
    const int warp_id = tid / 32;
50
    const int warp_row = warp_id / 2;
51
    const int warp_col_group = warp_id % 2;
52

53
    wmma::fragment<wmma::accumulator, 16, 16, WMMA_K, float>
54
        accumulators[OUTPUT_FRAGMENTS_PER_WARP];
55
    #pragma unroll
56
    for (int fragment = 0; fragment < OUTPUT_FRAGMENTS_PER_WARP; ++fragment) {
57
        wmma::fill_fragment(accumulators[fragment], 0.0f);
58
    }
59

60
    for (int inner_base = 0; inner_base < N; inner_base += BK) {
61
        for (int index = tid; index < BM * BK;
62
             index += THREADS_PER_BLOCK) {
63
            const int local_row = index / BK;
64
            const int local_inner = index % BK;
65
            const int row = blockIdx.y * BM + local_row;
66
            const float value = A[static_cast<size_t>(row) * N +
67
                                  inner_base + local_inner];
68
            shared_a[local_row][local_inner] = to_tf32(value);
69
        }
70
        for (int index = tid; index < BK * BN;
71
             index += THREADS_PER_BLOCK) {
72
            const int local_inner = index / BN;
73
            const int local_col = index % BN;
74
            const int col = blockIdx.x * BN + local_col;
75
            const float value = B[static_cast<size_t>(inner_base + local_inner) * K +
76
                                  col];
77
            shared_b[local_inner][local_col] = to_tf32(value);
78
        }
79
        __syncthreads();
80

81
        const int local_row = warp_row * 16;
82
        #pragma unroll
83
        for (int inner = 0; inner < BK; inner += WMMA_K) {
84
            wmma::fragment<wmma::matrix_a, 16, 16, WMMA_K,
85
                           wmma::precision::tf32, wmma::row_major> a_fragment;
86
            wmma::load_matrix_sync(
87
                a_fragment, &shared_a[local_row][inner], BK);
88

89
            #pragma unroll
90
            for (int fragment = 0;
91
                 fragment < OUTPUT_FRAGMENTS_PER_WARP;
92
                 ++fragment) {
93
                wmma::fragment<wmma::matrix_b, 16, 16, WMMA_K,
94
                               wmma::precision::tf32,
95
                               wmma::row_major> b_fragment;
96
                wmma::load_matrix_sync(
97
                    b_fragment,
98
                    &shared_b[inner][warp_col_group * 64 + fragment * 16],
99
                    BN);
100
                wmma::mma_sync(accumulators[fragment], a_fragment,
101
                               b_fragment, accumulators[fragment]);
102
            }
103
        }
104
        __syncthreads();
105
    }
106

107
    const int output_row = blockIdx.y * BM + warp_row * 16;
108
    const int output_col = blockIdx.x * BN + warp_col_group * 64;
109
    #pragma unroll
110
    for (int fragment = 0; fragment < OUTPUT_FRAGMENTS_PER_WARP; ++fragment) {
111
        wmma::store_matrix_sync(
112
            &C[static_cast<size_t>(output_row) * K +
113
               output_col + fragment * 16],
114
            accumulators[fragment], K, wmma::mem_row_major);
115
    }
116
#endif
117
}
118

119
}  // namespace
120

121
extern "C" void solve(const float* A, const float* B, float* C,
122
                      int M, int N, int K) {
123
    static const bool supports_wmma_tf32 = [] {
124
        int device = 0;
125
        cudaDeviceProp properties{};
126
        cudaGetDevice(&device);
127
        cudaGetDeviceProperties(&properties, device);
128
        return properties.major >= 8;
129
    }();
130

131
    if (!supports_wmma_tf32 ||
132
        M % BM != 0 || N % BK != 0 || K % BN != 0) {
133
        const dim3 block(16, 16);
134
        const dim3 grid((K + 15) / 16, (M + 15) / 16);
135
        scalar_fallback_v5<<<grid, block>>>(A, B, C, M, N, K);
136
        return;
137
    }
138

139
    const dim3 block(THREADS_PER_BLOCK);
140
    const dim3 grid(K / BN, M / BM);
141
    matrix_multiplication_v5_wmma_tf32<<<grid, block>>>(A, B, C, M, N, K);
142
}

展开收起

Test Methodology#

所有 CUDA 方法都使用同一个测试文件，并且必须通过 9 组正确性 shape，其中包括标量、矩形、对齐、未对齐、partial-tile 和 multi-block 维度。检查器会拒绝非有限输出值，并使用 5e-3 的缩放容差将每个输出元素与 CPU reference 比较。该容差覆盖 V5 有意进行的 TF32 输入舍入；严格的 FP32 方法也使用同一测试评估，从而确保对比中的每一行都遵循统一流程。129 x 96 乘 96 x 132 的 case 用于覆盖 V4 处理不完整输出 tile 的 asynchronous fast path，而 256 x 96 乘 96 x 256 的 case 则用于覆盖多个 WMMA output block。

性能测试 case 使用 M=8192、N=6144 和 K=4096。输入被初始化为确定性的非零值，以免全零分配产生缺乏代表性的内存行为。每种方法都针对远程 GPU 的原生架构编译。在 10 次基于 CUDA event 的测量前，先执行 5 次 warm-up iteration。下表给出这 10 次测量的平均值、最小值、最大值和计算得到的浮点吞吐率。

Test Code#

1
#include <cuda_runtime.h>
2

3
#include <array>
4
#include <cmath>
5
#include <iostream>
6

7
extern "C" void solve(const float* a, const float* b, float* c,
8
                      int m, int n, int k);
9

10
int main() {
11
    constexpr int m = 2;
12
    constexpr int n = 2;
13
    constexpr int k = 2;
14
    const std::array<float, 4> matrix_a = {1.0f, 2.0f, 3.0f, 4.0f};
15
    const std::array<float, 4> matrix_b = {5.0f, 6.0f, 7.0f, 8.0f};
16
    const std::array<float, 4> expected = {19.0f, 22.0f, 43.0f, 50.0f};
17
    std::array<float, 4> actual = {};
18
    constexpr size_t bytes = 4 * sizeof(float);
19

20
    float* device_a = nullptr;
21
    float* device_b = nullptr;
22
    float* device_c = nullptr;
23

24
    auto check_cuda = [](cudaError_t status, const char* operation) {
25
        if (status == cudaSuccess) {
26
            return true;
27
        }
28
        std::cerr << "Test failed. CUDA error in " << operation << ": "
29
                  << cudaGetErrorString(status) << '\n';
30
        return false;
31
    };
32

33
    bool success = check_cuda(cudaMalloc(&device_a, bytes), "cudaMalloc(device_a)") &&
34
                   check_cuda(cudaMalloc(&device_b, bytes), "cudaMalloc(device_b)") &&
35
                   check_cuda(cudaMalloc(&device_c, bytes), "cudaMalloc(device_c)") &&
36
                   check_cuda(cudaMemcpy(device_a, matrix_a.data(), bytes,
37
                                         cudaMemcpyHostToDevice), "copy matrix_a") &&
38
                   check_cuda(cudaMemcpy(device_b, matrix_b.data(), bytes,
39
                                         cudaMemcpyHostToDevice), "copy matrix_b");
40

41
    if (success) {
42
        solve(device_a, device_b, device_c, m, n, k);
43
        success = check_cuda(cudaGetLastError(), "solve") &&
44
                  check_cuda(cudaDeviceSynchronize(), "synchronize") &&
45
                  check_cuda(cudaMemcpy(actual.data(), device_c, bytes,
46
                                        cudaMemcpyDeviceToHost), "copy result");
47
    }
48

49
    cudaFree(device_a);
50
    cudaFree(device_b);
51
    cudaFree(device_c);
52

53
    if (!success) {
54
        return 1;
55
    }
56

57
    for (size_t index = 0; index < actual.size(); ++index) {
58
        if (std::fabs(actual[index] - expected[index]) > 1e-5f) {
59
            std::cerr << "Test failed at index " << index
60
                      << ". Expected: " << expected[index]
61
                      << ", Actual: " << actual[index] << '\n';
62
            return 1;
63
        }
64
    }
65

66
    std::cout << "Test passed.\n";
67
    return 0;
68
}

展开收起

Test Result#

以下表格特意将同一款 GPU 的结果集中展示。它们用于比较同一平台上的不同方法，而不是为两款 GPU 排名。Default 当前选择 V3 large-tile 实现，因此该行是对同一源码路径的重复运行，并非另一种算法。

NVIDIA GeForce RTX 5070 Ti Laptop GPU#

Method	Arithmetic	Status	Average Time	Minimum Time	Maximum Time	Performance
V0 naive	FP32	PASS	383.370 ms	348.679 ms	417.879 ms	1075.506 GFLOPS
V1 shared memory	FP32	PASS	250.884 ms	224.301 ms	302.463 ms	1643.454 GFLOPS
V1 1D thread tiling	FP32	PASS	135.891 ms	128.005 ms	150.389 ms	3034.181 GFLOPS
V1 2D thread tiling	FP32	PASS	123.423 ms	117.553 ms	135.046 ms	3340.669 GFLOPS
V2 vectorized, `BK=32`	FP32	PASS	46.376 ms	44.035 ms	48.737 ms	8890.723 GFLOPS
V2 vectorized, `BK=16`	FP32	PASS	38.761 ms	36.343 ms	40.281 ms	10637.490 GFLOPS
V2 vectorized, `128 x 128`	FP32	PASS	38.842 ms	37.423 ms	41.546 ms	10615.226 GFLOPS
V3 software double buffer, `64 x 64`	FP32	PASS	39.642 ms	34.580 ms	42.673 ms	10400.993 GFLOPS
V3 software double buffer, `128 x 128`	FP32	PASS	34.590 ms	32.101 ms	37.409 ms	11920.091 GFLOPS
V4 hardware `cp.async`	FP32	PASS	34.572 ms	31.259 ms	37.045 ms	11926.402 GFLOPS
V5 WMMA Tensor Core	TF32 input, FP32 accumulation	PASS	57.805 ms	52.967 ms	66.704 ms	7132.925 GFLOPS
Default (V3 `128 x 128`)	FP32	PASS	34.486 ms	31.874 ms	37.132 ms	11956.000 GFLOPS

在这款 GPU 上，V4 在最终统一测试中的平均时间最低，但考虑到观察到的运行间波动，其 34.572 ms 结果实际上与匹配的 large V3 的 34.590 ms 持平。V4 多次运行的时间大约在 32 至 35 ms 之间。因此，证据表明 hardware-asynchronous copy 在该平台上具有竞争力，而不是 0.018 ms 的最终差异具有显著意义。V5 在文档所述的 TF32 容差下结果正确，但它慢于经过调优的 CUDA Core kernel，因为这个手写 WMMA 映射需要承担大量转换、shared-memory、同步和 fragment 管理开销。

NVIDIA GeForce RTX 4090#

Method	Arithmetic	Status	Average Time	Minimum Time	Maximum Time	Performance
V0 naive	FP32	PASS	83.140 ms	82.831 ms	83.375 ms	4959.321 GFLOPS
V1 shared memory	FP32	PASS	59.379 ms	58.985 ms	59.865 ms	6943.818 GFLOPS
V1 1D thread tiling	FP32	PASS	19.484 ms	19.393 ms	19.505 ms	21161.508 GFLOPS
V1 2D thread tiling	FP32	PASS	13.446 ms	13.420 ms	13.465 ms	30665.525 GFLOPS
V2 vectorized, `BK=32`	FP32	PASS	11.646 ms	11.098 ms	12.106 ms	35404.086 GFLOPS
V2 vectorized, `BK=16`	FP32	PASS	11.314 ms	11.213 ms	11.345 ms	36441.761 GFLOPS
V2 vectorized, `128 x 128`	FP32	PASS	13.146 ms	13.135 ms	13.173 ms	31363.839 GFLOPS
V3 software double buffer, `64 x 64`	FP32	PASS	11.514 ms	11.380 ms	11.654 ms	35809.707 GFLOPS
V3 software double buffer, `128 x 128`	FP32	PASS	10.052 ms	9.939 ms	10.213 ms	41019.869 GFLOPS
V4 hardware `cp.async`	FP32	PASS	10.293 ms	10.230 ms	10.415 ms	40056.186 GFLOPS
V5 WMMA Tensor Core	TF32 input, FP32 accumulation	PASS	12.675 ms	12.107 ms	13.133 ms	32528.826 GFLOPS
Default (V3 `128 x 128`)	FP32	PASS	10.096 ms	10.010 ms	10.144 ms	40840.572 GFLOPS

在这款 GPU 上，large V3 软件流水线仍是实测最快的实现。V4 将复制与算术计算重叠执行，但未能抵消额外的 copy-group 和同步开销。V5 再次证明 Tensor Core 执行功能正常，却没有超越已有的 FP32 路径。这一结果说明了一个有用的边界：当数据准备和 fragment 调度尚未达到生产级 GEMM 库的成熟程度时，仅凭 Tensor Core 的峰值吞吐率并不能保证端到端 kernel 更快。

Triton#

Approach#

Triton 实现计算与 CUDA 相同的行主序矩阵乘积：

C[m,k]=∑n=0N−1A[m,n]B[n,k].C[m,k] = \sum_{n=0}^{N-1} A[m,n]B[n,k].C[m,k]=n=0∑N−1A[m,n]B[n,k].

一个 Triton program instance 负责一个 64 x 64 输出 tile。与 CUDA 源码不同，该实现不会将单个标量输出显式分配给 threadIdx 坐标，而是描述完整的 index block 和 value block；Triton 编译器再将这些操作映射到 GPU thread 和 warp。reduction 维度以 32 为单位分块处理：

1
BLOCK_SIZE_M = 64
2
BLOCK_SIZE_N = 32
3
BLOCK_SIZE_K = 64
4
GROUP_SIZE_M = 8

输出 grid 被展平为一维。tl.program_id(0) 用于标识当前 program instance。kernel 根据这个 linear ID 推导 PID_M 和 PID_K，二者分别标识输出行 tile 和输出列 tile。完整 grid 包含：

⌈M64⌉×⌈K64⌉\left\lceil \frac{M}{64} \right\rceil \times \left\lceil \frac{K}{64} \right\rceil⌈64M⌉×⌈64K⌉

个 program instance。

program 按最多 8 个 M tile 分组排序，之后才沿 K 继续推进。这只会改变执行顺序，不会改变每个 program 所负责的输出 tile。相邻 program 更有可能复用 A 和 B 的 cache 区域；相比之下，简单的行主序排列可能会沿一个输出维度前进较远后，才重新访问可复用的输入数据。

对于一个 program，行列 offset 分别为：

1
a_m_offset = PID_M * 64 + [0, ..., 63]
2
b_k_offset = PID_K * 64 + [0, ..., 63]

在每个 reduction step 中，tl.arange 创建行、reduction 和列 offset。使用 [:, None] 和 [None, :] 添加 singleton dimension 后，这些 vector 会被广播为来自 A 的 64 x 32 pointer tile 和来自 B 的 32 x 64 pointer tile。A、B 和 C 的行主序 stride 分别以 (N, 1)、(K, 1) 和 (K, 1) 显式传入。

mask 在正确性方面所起的作用，与 CUDA kernel 中的边界检查和零填充相同。被 mask 的输入位置会按零加载，因此不会对 reduction 产生贡献；被 mask 的输出位置则不会被存储。因此，同一个 kernel 即可处理不是 64 或 32 的倍数的维度，无需启动单独的边界 kernel。

tl.dot 将两个输入 tile 相乘，并累加到一个 float32 64 x 64 block 中。allow_tf32=False 使计算保持在该实现采用的 non-TF32 路径上。reduction loop 结束后，最终的 mask 会防止向 C 范围外存储。

launch 为每个 program 使用 4 个 warp 和 3 个 pipeline stage。与此前固定的 32 x 32 输出 tile 和未分组的二维 grid 相比，更大的 tile 使每个 program 完成更多工作，分组排序则改善了局部性。在统一 benchmark 下，这使 RTX 4090 上的平均时间从 16.907 ms 降至 9.803 ms，并使 RTX 5070 Ti Laptop GPU 上的平均时间从 78.710 ms 降至 32.518 ms。

Solution#

1
import torch
2
import triton
3
import triton.language as tl
4

5

6
@triton.jit
7
def matrix_multiplication_kernel(
8
    a, b, c, M, N, K, stride_am, stride_an, stride_bn, stride_bk, stride_cm, stride_ck,
9
    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr,
10
    BLOCK_SIZE_K: tl.constexpr, GROUP_SIZE_M: tl.constexpr
11
):
12
    pid = tl.program_id(0)
13
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
14
    num_pid_k = tl.cdiv(K, BLOCK_SIZE_K)
15
    num_pid_in_group = GROUP_SIZE_M * num_pid_k
16
    group_id = pid // num_pid_in_group
17
    first_pid_m = group_id * GROUP_SIZE_M
18
    group_size_m = tl.minimum(num_pid_m - first_pid_m, GROUP_SIZE_M)
19
    PID_M = first_pid_m + (pid % num_pid_in_group) % group_size_m
20
    PID_K = (pid % num_pid_in_group) // group_size_m
21
    MAX_N = tl.cdiv(N, BLOCK_SIZE_N)
22

23
    accumulated_block = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_K), dtype=tl.float32)
24

25
    start_a_m = PID_M * BLOCK_SIZE_M
26
    a_m_offset = start_a_m + tl.arange(0, BLOCK_SIZE_M)
27

28
    start_b_k = PID_K * BLOCK_SIZE_K
29
    b_k_offset = start_b_k + tl.arange(0, BLOCK_SIZE_K)
30

31
    for n in tl.range(MAX_N):
32
        start_a_n = n * BLOCK_SIZE_N
33
        a_n_offset = start_a_n + tl.arange(0, BLOCK_SIZE_N)
34

35
        a_mn_mask = (a_m_offset[:, None] < M) & (a_n_offset[None, :] < N)
36
        a_mn_ptrs = a + a_m_offset[:, None] * stride_am + a_n_offset[None, :] * stride_an
37
        block_a_mn = tl.load(a_mn_ptrs, mask=a_mn_mask, other=0.0)
38

39
        start_b_n = n * BLOCK_SIZE_N
40
        b_n_offset = start_b_n + tl.arange(0, BLOCK_SIZE_N)
41

42

43
        b_nk_mask = (b_n_offset[:, None] < N) & (b_k_offset[None, :] < K)
44
        b_nk_ptrs = b + b_n_offset[:, None] * stride_bn + b_k_offset[None, :] * stride_bk
45

46
        block_b_nk = tl.load(b_nk_ptrs, mask=b_nk_mask, other=0.0)
47

48
        accumulated_block = tl.dot(block_a_mn, block_b_nk, accumulated_block, allow_tf32=False)
49
        # block_ab = tl.dot(block_a_mn, block_b_nk)
50
        # accumulated_block += block_ab
51

52
    c_mk_ptrs = c + a_m_offset[:, None] * stride_cm + b_k_offset[None, :] * stride_ck
53
    c_mk_mask = (a_m_offset[:, None] < M) & (b_k_offset[None, :] < K)
54
    tl.store(c_mk_ptrs, accumulated_block, mask=c_mk_mask)
55

56
# a, b, c are tensors on the GPU
57
def solve(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, M: int, N: int, K: int):
58
    stride_am, stride_an = N, 1
59
    stride_bn, stride_bk = K, 1
60
    stride_cm, stride_ck = K, 1
61

62
    BLOCK_SIZE_M = 64
63
    BLOCK_SIZE_N = 32
64
    BLOCK_SIZE_K = 64
65
    GROUP_SIZE_M = 8
66

67
    grid = (triton.cdiv(M, BLOCK_SIZE_M) * triton.cdiv(K, BLOCK_SIZE_K),)
68
    matrix_multiplication_kernel[grid](
69
        a, b, c, M, N, K, stride_am, stride_an, stride_bn, stride_bk, stride_cm, stride_ck,
70
        BLOCK_SIZE_M = BLOCK_SIZE_M,
71
        BLOCK_SIZE_N = BLOCK_SIZE_N,
72
        BLOCK_SIZE_K = BLOCK_SIZE_K,
73
        GROUP_SIZE_M = GROUP_SIZE_M,
74
        num_warps=4,
75
        num_stages=3,
76
    )

展开收起

Test Code#

1
import importlib.util
2
import sys
3
from pathlib import Path
4

5
import torch
6

7

8
def load_implementation():
9
    source_path = Path(__file__).resolve().parents[2] / "src" / "triton" / "02_Matrix_Multiplication.py"
10
    spec = importlib.util.spec_from_file_location("matrix_multiplication_triton_minimal", source_path)
11
    if spec is None or spec.loader is None:
12
        raise RuntimeError(f"Unable to load implementation: {source_path}")
13
    module = importlib.util.module_from_spec(spec)
14
    spec.loader.exec_module(module)
15
    return module
16

17

18
def main():
19
    matrix_a = torch.tensor([[1.0, 2.0], [3.0, 4.0]], device="cuda")
20
    matrix_b = torch.tensor([[5.0, 6.0], [7.0, 8.0]], device="cuda")
21
    expected = torch.tensor([[19.0, 22.0], [43.0, 50.0]], device="cuda")
22
    actual = torch.empty_like(expected)
23

24
    load_implementation().solve(matrix_a, matrix_b, actual, 2, 2, 2)
25
    torch.cuda.synchronize()
26
    torch.testing.assert_close(actual, expected)
27
    print("Test passed.")
28

29

30
if __name__ == "__main__":
31
    try:
32
        main()
33
    except Exception as error:
34
        print(f"Test failed: {error}", file=sys.stderr)
35
        raise

展开收起

Test Result#

Platform	Status	Problem Size	Iterations	Average Time	Minimum Time	Maximum Time	Performance
NVIDIA GeForce RTX 5070 Ti Laptop GPU	PASS	M=8192, N=6144, K=4096	10	32.518 ms	30.066 ms	34.157 ms	12679.673 GFLOPS
NVIDIA GeForce RTX 4090	PASS	M=8192, N=6144, K=4096	10	9.803 ms	8.994 ms	10.005 ms	42061.461 GFLOPS

PyTorch#

Approach#

PyTorch 实现刻意保持简短，因为该框架已经以内置 tensor operation 的形式提供了矩阵乘法：

1
torch.matmul(A, B, out=C)

在本题中，A 和 B 是 shape 分别为 M x N 和 N x K 的二维 tensor。torch.matmul 会检查内部维度是否兼容，计算 M x K 的矩阵乘积，并将结果直接写入传入的输出 tensor C。使用 out=C 很重要，因为题目要求最终结果必须存放在调用方提供的输出 tensor 中。

PyTorch 提供了多种乘法接口，但它们的含义并不完全相同：

A * B 执行逐元素乘法，并非矩阵乘法。
A @ B 是矩阵乘法的 Python 运算符语法，遵循 torch.matmul 语义。
torch.mm(A, B) 专用于二维矩阵输入。
torch.bmm(A, B) 用于处理成批的三维矩阵，不支持通用 broadcasting。
torch.matmul(A, B) 支持 vector、matrix、更高维 batch 和 broadcasting。
torch.einsum(...) 通过 index notation 显式表达 contraction，适用于更复杂的 tensor 关系。

torch.mm 和 A @ B 都可以表达本题的二维矩阵运算，但实际源码使用的是带 out argument 的 torch.matmul，因此本文档说明和测试的也是这一操作。

该实现中没有显式的 grid、thread index、tile size、mask 或同步。PyTorch 根据 tensor 的 device 和 dtype 分派操作，并由其 CUDA backend 选择底层 GPU 实现。所需的 solve signature 仍保留 M、N 和 K argument，但不会直接读取它们；实际执行的矩阵乘法由 tensor shape 决定。

Solution#

1
import torch
2

3
# A, B, C are tensors on the GPU
4
def solve(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, M: int, N: int, K: int):
5
    torch.matmul(A,B,out=C)

Test Code#

1
import importlib.util
2
import sys
3
from pathlib import Path
4

5
import torch
6

7

8
def load_implementation():
9
    source_path = Path(__file__).resolve().parents[2] / "src" / "pytorch" / "02_Matrix_Multiplication.py"
10
    spec = importlib.util.spec_from_file_location("matrix_multiplication_pytorch_minimal", source_path)
11
    if spec is None or spec.loader is None:
12
        raise RuntimeError(f"Unable to load implementation: {source_path}")
13
    module = importlib.util.module_from_spec(spec)
14
    spec.loader.exec_module(module)
15
    return module
16

17

18
def main():
19
    matrix_a = torch.tensor([[1.0, 2.0], [3.0, 4.0]], device="cuda")
20
    matrix_b = torch.tensor([[5.0, 6.0], [7.0, 8.0]], device="cuda")
21
    expected = torch.tensor([[19.0, 22.0], [43.0, 50.0]], device="cuda")
22
    actual = torch.empty_like(expected)
23

24
    load_implementation().solve(matrix_a, matrix_b, actual, 2, 2, 2)
25
    torch.cuda.synchronize()
26
    torch.testing.assert_close(actual, expected)
27
    print("Test passed.")
28

29

30
if __name__ == "__main__":
31
    try:
32
        main()
33
    except Exception as error:
34
        print(f"Test failed: {error}", file=sys.stderr)
35
        raise

展开收起

Test Result#

Platform	Status	Problem Size	Iterations	Average Time	Minimum Time	Maximum Time	Performance
NVIDIA GeForce RTX 5070 Ti Laptop GPU	PASS	M=8192, N=6144, K=4096	10	29.305 ms	25.729 ms	31.689 ms	14069.752 GFLOPS
NVIDIA GeForce RTX 4090	PASS	M=8192, N=6144, K=4096	10	7.469 ms	7.188 ms	7.559 ms	55204.880 GFLOPS

References#

Acknowledgements#

本文档参考了 Du Ziyuan 的文章 CUDA Learning Journey [11] — A Detailed Explanation of Matrix Multiplication 中的讲解材料、图示和动画。原作采用 CC BY-NC-SA 4.0 许可；改编后的文档和在本地重新制作的媒体内容也按相同许可发布。改动包括围绕 LeetGPU 接口重新组织讲解，使讨论和代码与本项目的实现相匹配，并加入项目专用测试和实测结果。

关注 LeetGPU repository，了解更多内容。

LeetGPU | 01_Vector_Addition

Mon, 20 Jul 2026 00:00:00 GMT

关注 LeetGPU repository，了解更多内容。

Problem Description#

LeetGPU Challenge

Write a GPU program that performs element-wise addition of two vectors containing 32-bit floating point numbers. The program should take two input vectors of equal length and produce a single output vector containing their sum.

Implementation Requirements#

External libraries are not permitted.
The solve function signature must remain unchanged.
The final result must be stored in vector C.

Example 1#

1
Input:  A = [1.0, 2.0, 3.0, 4.0]
2
        B = [5.0, 6.0, 7.0, 8.0]
3
Output: C = [6.0, 8.0, 10.0, 12.0]

Example 2#

1
Input:  A = [1.5, 1.5, 1.5]
2
        B = [2.3, 2.3, 2.3]
3
Output: C = [3.8, 3.8, 3.8]

Constraints#

Input vectors A and B have identical lengths.
1 <= N <= 100,000,000.
Performance is measured with N = 25,000,000.

CUDA#

Approach#

CUDA 采用“一元素一线程”的方式，为每个向量元素分配一个 GPU 线程。

线程的全局索引由 blockIdx.x、blockDim.x 和 threadIdx.x 共同计算得出。
每个索引有效的线程分别从 A 和 B 读取一个元素，完成加法后将结果写入 C。
每个线程块包含 256 个线程，线程块数量通过向上取整除法确定。
idx < N 的边界检查可避免最后一个线程块发生越界访问。
内核使用 CUDA C++ 编写，并通过 nvcc 编译。

Solution#

1
#include <cuda_runtime.h>
2

3
__global__ void vector_add(const float* A, const float* B, float* C, int N) {
4
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
5
    if (idx < N) {
6
        C[idx] = A[idx] + B[idx];
7
    }
8
}
9

10
// A, B, C are device pointers (i.e. pointers to memory on the GPU)
11
extern "C" void solve(const float* A, const float* B, float* C, int N) {
12
    int threadsPerBlock = 256;
13
    int blocksPerGrid = (N + threadsPerBlock - 1) / threadsPerBlock;
14

15
    vector_add<<<blocksPerGrid, threadsPerBlock>>>(A, B, C, N);
16
}

展开收起

Test Code#

1
#include <cuda_runtime.h>
2

3
#include <array>
4
#include <cmath>
5
#include <iostream>
6

7
extern "C" void solve(const float* a, const float* b, float* c, int n);
8

9
int main() {
10
    constexpr int n = 4;
11
    const std::array<float, n> input_a = {1.0f, 2.0f, 3.0f, 4.0f};
12
    const std::array<float, n> input_b = {5.0f, 6.0f, 7.0f, 8.0f};
13
    const std::array<float, n> expected = {6.0f, 8.0f, 10.0f, 12.0f};
14
    std::array<float, n> actual = {};
15
    constexpr size_t bytes = n * sizeof(float);
16

17
    float* device_a = nullptr;
18
    float* device_b = nullptr;
19
    float* device_c = nullptr;
20

21
    auto check_cuda = [](cudaError_t status, const char* operation) {
22
        if (status == cudaSuccess) {
23
            return true;
24
        }
25
        std::cerr << "Test failed. CUDA error in " << operation << ": "
26
                  << cudaGetErrorString(status) << '\n';
27
        return false;
28
    };
29

30
    bool success = check_cuda(cudaMalloc(&device_a, bytes), "cudaMalloc(device_a)") &&
31
                   check_cuda(cudaMalloc(&device_b, bytes), "cudaMalloc(device_b)") &&
32
                   check_cuda(cudaMalloc(&device_c, bytes), "cudaMalloc(device_c)") &&
33
                   check_cuda(cudaMemcpy(device_a, input_a.data(), bytes, cudaMemcpyHostToDevice), "copy input_a") &&
34
                   check_cuda(cudaMemcpy(device_b, input_b.data(), bytes, cudaMemcpyHostToDevice), "copy input_b");
35

36
    if (success) {
37
        solve(device_a, device_b, device_c, n);
38
        success = check_cuda(cudaGetLastError(), "solve") &&
39
                  check_cuda(cudaDeviceSynchronize(), "synchronize") &&
40
                  check_cuda(cudaMemcpy(actual.data(), device_c, bytes, cudaMemcpyDeviceToHost), "copy result");
41
    }
42

43
    cudaFree(device_a);
44
    cudaFree(device_b);
45
    cudaFree(device_c);
46

47
    if (!success) {
48
        return 1;
49
    }
50

51
    for (int index = 0; index < n; ++index) {
52
        if (std::fabs(actual[index] - expected[index]) > 1e-6f) {
53
            std::cerr << "Test failed at index " << index
54
                      << ". Expected: " << expected[index]
55
                      << ", Actual: " << actual[index] << '\n';
56
            return 1;
57
        }
58
    }
59

60
    std::cout << "Test passed.\n";
61
    return 0;
62
}

展开收起

Test Result#

Platform	Status	Problem Size	Iterations	Average Time	Minimum Time	Maximum Time	Performance
NVIDIA GeForce RTX 5070 Ti Laptop GPU	PASS	N=25000000	20	0.826 ms	0.748 ms	1.325 ms	363.144 GB/s
NVIDIA GeForce RTX 4090	PASS	N=25000000	20	0.324 ms	0.322 ms	0.327 ms	926.118 GB/s
NVIDIA GeForce RTX 5090	PASS	N=25000000	20	0.193 ms	0.190 ms	0.195 ms	1557.852 GB/s

Triton#

Approach#

Triton 让每个程序实例负责处理一段连续的向量元素。

tl.program_id(0) 用于标识当前程序实例。
每个程序实例最多处理由 tl.arange 生成的 1024 个元素。
内核分别加载两个数据块，逐元素相加后存储结果。
掩码可避免处理最后一个数据块时发生越界加载和存储。
内核使用 Python 编写，并由 Triton 即时编译（JIT）为 GPU 代码。

Solution#

1
import torch
2
import triton
3
import triton.language as tl
4

5

6
@triton.jit
7
def vector_add_kernel(a, b, c, n_elements, BLOCK_SIZE: tl.constexpr):
8
    pid = tl.program_id(0)
9

10
    block_start = pid * BLOCK_SIZE
11
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
12
    mask=offsets<n_elements
13

14
    a_block = tl.load(a+offsets, mask=mask)
15
    b_block = tl.load(b+offsets, mask=mask)
16

17
    c_block = a_block+b_block
18

19
    tl.store(c+offsets, c_block, mask=mask)
20

21

22
# a, b, c are tensors on the GPU
23
def solve(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, N: int):
24
    BLOCK_SIZE = 1024
25
    grid = (triton.cdiv(N, BLOCK_SIZE),)
26
    vector_add_kernel[grid](a, b, c, N, BLOCK_SIZE)

展开收起

Test Code#

1
import importlib.util
2
import sys
3
from pathlib import Path
4

5
import torch
6

7

8
def load_implementation():
9
    source_path = Path(__file__).resolve().parents[2] / "src" / "triton" / "01_Vector_Addition.py"
10
    spec = importlib.util.spec_from_file_location("vector_addition_triton_minimal", source_path)
11
    if spec is None or spec.loader is None:
12
        raise RuntimeError(f"Unable to load implementation: {source_path}")
13
    module = importlib.util.module_from_spec(spec)
14
    spec.loader.exec_module(module)
15
    return module
16

17

18
def main():
19
    input_a = torch.tensor([1.0, 2.0, 3.0, 4.0], device="cuda")
20
    input_b = torch.tensor([5.0, 6.0, 7.0, 8.0], device="cuda")
21
    expected = torch.tensor([6.0, 8.0, 10.0, 12.0], device="cuda")
22
    actual = torch.empty_like(input_a)
23

24
    load_implementation().solve(input_a, input_b, actual, input_a.numel())
25
    torch.cuda.synchronize()
26
    torch.testing.assert_close(actual, expected)
27
    print("Test passed.")
28

29

30
if __name__ == "__main__":
31
    try:
32
        main()
33
    except Exception as error:
34
        print(f"Test failed: {error}", file=sys.stderr)
35
        raise

展开收起

Test Result#

Platform	Status	Problem Size	Iterations	Average Time	Minimum Time	Maximum Time	Performance
NVIDIA GeForce RTX 5070 Ti Laptop GPU	PASS	N=25000000	20	0.809 ms	0.764 ms	1.085 ms	371.021 GB/s
NVIDIA GeForce RTX 4090	PASS	N=25000000	20	0.341 ms	0.338 ms	0.350 ms	880.608 GB/s
NVIDIA GeForce RTX 5090	PASS	N=25000000	20	0.223 ms	0.216 ms	0.251 ms	1344.240 GB/s

PyTorch#

Approach#

PyTorch 通过内置的 torch.add 算子完成向量加法。

torch.add(A, B, out=C) 将对应位置的元素相加，并把结果直接写入 C。
PyTorch 会自动选择合适的 GPU 内核并发起执行。
线程调度、索引计算和边界处理均由 PyTorch 在内部完成。
为满足接口要求，函数签名中保留了 N，但该实现不会直接使用它。

Solution#

1
import torch
2

3

4
# A, B, C are tensors on the GPU
5
def solve(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, N: int):
6
    torch.add(A,B,out = C);

Test Code#

1
import importlib.util
2
import sys
3
from pathlib import Path
4

5
import torch
6

7

8
def load_implementation():
9
    source_path = Path(__file__).resolve().parents[2] / "src" / "pytorch" / "01_Vector_Addition.py"
10
    spec = importlib.util.spec_from_file_location("vector_addition_pytorch_minimal", source_path)
11
    if spec is None or spec.loader is None:
12
        raise RuntimeError(f"Unable to load implementation: {source_path}")
13
    module = importlib.util.module_from_spec(spec)
14
    spec.loader.exec_module(module)
15
    return module
16

17

18
def main():
19
    input_a = torch.tensor([1.0, 2.0, 3.0, 4.0], device="cuda")
20
    input_b = torch.tensor([5.0, 6.0, 7.0, 8.0], device="cuda")
21
    expected = torch.tensor([6.0, 8.0, 10.0, 12.0], device="cuda")
22
    actual = torch.empty_like(input_a)
23

24
    load_implementation().solve(input_a, input_b, actual, input_a.numel())
25
    torch.cuda.synchronize()
26
    torch.testing.assert_close(actual, expected)
27
    print("Test passed.")
28

29

30
if __name__ == "__main__":
31
    try:
32
        main()
33
    except Exception as error:
34
        print(f"Test failed: {error}", file=sys.stderr)
35
        raise

展开收起

Test Result#

Platform	Status	Problem Size	Iterations	Average Time	Minimum Time	Maximum Time	Performance
NVIDIA GeForce RTX 5070 Ti Laptop GPU	PASS	N=25000000	20	0.817 ms	0.731 ms	1.077 ms	367.031 GB/s
NVIDIA GeForce RTX 4090	PASS	N=25000000	20	0.330 ms	0.329 ms	0.334 ms	909.717 GB/s
NVIDIA GeForce RTX 5090	PASS	N=25000000	20	0.207 ms	0.204 ms	0.216 ms	1450.363 GB/s

References#

LeetGPU Challenge

关注 LeetGPU repository，了解更多内容。