Denoising Diffusion Implicit Models

Posted on 2025-05-14 Edited on 2025-05-15 In Diffusion Valine:

URL

paper: https://arxiv.org/pdf/2010.02502

TL;DR

这篇论文提出了去噪扩散隐式模型 Denoising Diffusion Implicit Models (DDIM) 模型，可以看作是对 Denoising Diffusion Probabilistic Models (DDPM) 模型的改进。
DDPM 的采样过程是一个 Markov 过程，Markov 过程只有知道 t 时刻的状态才能计算第 t-1 时刻，而 DDIM 的采样过程是一个非 Markov 过程。
DDIM 的优势是：
1. 去噪速度更快。可以在采样（去噪）时使用更少的时间步数，从而加快采样速度，并且用 DDPM 训练的模型可以直接用于 DDIM 的采样，二者可以无缝衔接。
2. 确定性。在 DDIM 中，给定一个模型和一个噪声图像和时间步数，可以确定性地生成一个图像（运行再多次也是同一个图）。在 DDPM 中，给定一个模型和一个噪声图像和时间步数，生成的图像是随机的（每次跑都不一样）。

Algorithm

DDIM 的公式是从 DDPM 公式通过复杂推导得到的，推导过程比较复杂，这里不做详细介绍，重点讲二者逆向过程（去噪）公式区别和使用上的区别。

DDPM 逆向过程公式

x_{t-1}=\frac{1}{\sqrt{\alpha_t}}(x_t-\frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\cdot \epsilon_\theta(x_t,t))

其中：
- $\epsilon_\theta(x_t,t)$ 是模型预测的噪声，即 model(x_t, t)
- $\alpha_t=1-\beta_t$
- $\bar\alpha_t=\prod_{s=1}^t\alpha_s$
- t 取值是 999, 998,..., 1, 0，即从 T-1 到 0 逐步去噪

DDIM 逆向过程公式

x_{t-1}=\sqrt {\bar\alpha_t-1}(\frac{x_t-\sqrt {\bar\alpha_t-1}\cdot\epsilon_\theta(x_t,t)}{\sqrt {\bar\alpha_t}})+\sqrt{1-\bar\alpha_{t-1}}\cdot\epsilon_\theta(x_t,t)

其中：
- 大多数符号含义都和 DDPM 一样
- 只有 t 取值是 980, 960,..., 20, 0 这种非连续的取值

`DDIM` 公式拆解：

1. 预测原输入

pred_{x0}=\frac{x_t-\sqrt {\bar\alpha_t-1}\cdot\epsilon_\theta(x_t,t)}{\sqrt{\bar\alpha_t}}

通过当前噪声隐变量 $x_t$ 和模型预测的噪声 $\epsilon_\theta(x_t,t)$ 估计原始输入 $x_0$ ，即去噪

2. 计算调整方向

direction\_point=\sqrt{1-\bar\alpha_{t-1}}\cdot\epsilon_\theta(x_t,t)

根据噪声预测结果，计算从当前时间步 t 到前一个时间步 t-1 的调整方向，这一方向结合了噪声预测和噪声调度参数，用于引导隐变量的更新。

3. 更新隐变量

x_{t-1}=\sqrt {\bar\alpha_t-1}\cdot pred_{x0}+direction\_point

将预测的原始输入 $pred_{x0}$ 与调整方向结合，生成前一时刻的隐变量 $x_{t-1}$ 。此步骤通过线性组合逐步去噪，最终逼近目标数据 $x_0$ 。

DDPM 和 DDIM 使用上的区别

DDPM 去噪过程

model = UNet2DModel.from_pretrained("google/ddpm-celebahq-256").to(device)
scheduler = DDPMScheduler.from_pretrained("google/ddpm-celebahq-256")

# Get precalculated alphas and alpha bars from the scheduler
alphas = scheduler.alphas
alphas_cumprod = scheduler.alphas_cumprod

# Initialize sample with static random noise
sample = torch.load("random_noise.pt").to(device)

# DDPM denoising loop
# scheduler.timesteps = [999, 998, ..., 1, 0]
for t in tqdm.tqdm(scheduler.timesteps):
    with torch.no_grad():
        # Model prediction (noise residual)
        residual = model(sample, t).sample

        # DDPM denoising formula
        sample = (
            sample - (1 - alphas[t]) / torch.sqrt(1 - alphas_cumprod[t]) * residual
        ) / torch.sqrt(alphas[t])

        # Add random noise only for t > 1
        if t > 1:
            noise = torch.randn_like(sample).to(device)
            sample += torch.sqrt(1 - alphas[t]) * noise

DDIM 去噪过程

model = UNet2DModel.from_pretrained("google/ddpm-celebahq-256").to(device)
scheduler = DDIMScheduler.from_pretrained("google/ddpm-celebahq-256")
# set inference steps
scheduler.set_timesteps(num_inference_steps=50)

# Initialize sample with static random noise
sample = torch.load("random_noise.pt").to(device)

# DDIM denoising loop
# scheduler.timesteps = [980, 960,..., 20, 0]
for i, t in enumerate(tqdm.tqdm(scheduler.timesteps)):
    # 将时间步转换为LongTensor并确保在正确设备上
    t = t.to(device).long()

    # 获取当前和上一步的alpha累积乘积
    alpha_cumprod_t = scheduler.alphas_cumprod[t]
    alpha_cumprod_prev = (
        scheduler.alphas_cumprod[scheduler.timesteps[i + 1]]
        if i + 1 < len(scheduler.timesteps)
        else torch.tensor(1.0)
    )

    # 将alpha值转换到相同设备
    alpha_cumprod_t = alpha_cumprod_t.to(device)
    alpha_cumprod_prev = alpha_cumprod_prev.to(device)

    with torch.no_grad():
        # 1. 预测噪声残差
        residual = model(sample, t).sample

        # 2. 计算预测的原始图像x0（去噪后的图像）
        pred_x0 = (sample - torch.sqrt(1.0 - alpha_cumprod_t) * residual) / torch.sqrt(
            alpha_cumprod_t
        )

        # 3. 计算下一步的样本方向
        direction_xt = torch.sqrt(1.0 - alpha_cumprod_prev) * residual

        # 4. 组合得到新的样本
        sample = torch.sqrt(alpha_cumprod_prev) * pred_x0 + direction_xt

二者对比分析

DDIM 只需要 50 次迭代就能生成高质量的图像，而 DDPM 需要 1000 次迭代。
生成的图像质量相似，DDIM 生成的图像质量略高。
上面的代码中加载的噪声图是静态的，DDIM 跑多次生成的图像是一样的，而 DDPM 跑多次生成的图像是不一样的。
二者去噪结果对比，左侧是 DDIM，右侧是 DDPM：

Thoughts

DDIM 解决了 DDPM 的两大痛点，算是一个很好的改进。
为后续的 LDM 等模型打下了基础。

High-Resolution Image Synthesis with Latent Diffusion Models

Posted on 2025-05-11 In stable diffusion Valine:

URL

paper: https://arxiv.org/pdf/2112.10752
code: https://github.com/CompVis/latent-diffusion

TL;DR

Qwen3 技术报告先导篇

Posted on 2025-05-05 Edited on 2025-05-07 In LLM Valine:

URL

blog: https://qwenlm.github.io/zh/blog/qwen3/

TL;DR

Qwen3 系列模型四月二十九日正式发布，但到目前为止，还没发布技术报告，只有一篇官方博客，介绍了 Qwen3 的一些基本信息
本文围绕着官方博客介绍，结合实际使用情况，给出一些个人的理解

Qwen3 系列模型

本次 Qwen3 主要发布了八个版本的模型，其中包含两个 MOE 模型和六个 dense 模型：
- Qwen3-235B-A22B：MOE 模型，235B 参数量，22B 激活参数量
- Qwen3-30B-A3B：MOE 模型，30B 参数量，3B 激活参数量
- Qwen3-32B：dense 模型，32B 参数量
- Qwen3-14B：dense 模型，14B 参数量
- Qwen3-8B：dense 模型，8B 参数量
- Qwen3-4B：dense 模型，4B 参数量
- Qwen3-1.7B：dense 模型，1.7B 参数量
- Qwen3-0.6B：dense 模型，0.6B 参数量
还有对应的 Base 模型（只经过预训练）和 fp8 Quantized 模型（量化模型）
Qwen3 相较于上一代 Qwen2.5，一个较大的技术进步是：统一了 Reasoning 和非 Reasoning 模式，即不再区分推理模型和非推理模型（或者叫思考模型和非思考模型），而是一个模型可以通过 prompt 来选择推理模式和非推理模式

Qwen3 模型的主要特性

预训练

相较于上一代 Qwen2.5 使用了 18 万亿个 token 做预训练，Qwen3 使用了 36 万亿个 token 做预训练（整整翻了一倍，这是多大的数据团队才能搞出来的，~~太壕了~~）
包含了 119 种语言和方言，有大量数据是合成数据和在 PDF 上识别的文本数据
预训练分成了三个阶段：
- 第一阶段：模型在超过 30 万亿个 token 上进行了预训练，上下文长度为 4K token
- 第二阶段：通过增加知识密集型数据（如 STEM、编程和推理任务）的比例来改进数据集，在 5 万亿个 token 上进行了预训练
- 第三阶段：使用高质量的长上下文数据将上下文长度扩展到 32K token
预训练实际效果：官方的说法是 Qwen3-1.7B/4B/8B/14B/32B-Base 分别与 Qwen2.5-3B/7B/14B/32B/72B-Base 表现相当

后训练

简单概括：用 SFT → RL → SFT+RL 混合 → RL 四个阶段后训练 Qwen3-235B-A22B 和 Qwen3-32B 模型，然后蒸馏得到其他小尺寸模型（大 MOE 蒸小 MOE，大 dense 蒸小 dense）
四个阶段分别是：
- 长思维链冷启动：使用监督微调（SFT）训练模型生成初步的长思维链推理能力，作为初始阶段的基础。
- 长思维链强化学习：通过强化学习（RL）进一步提升模型在复杂推理任务中的表现，优化生成思维链的质量和连贯性。
- 思维模式融合：结合 SFT 和 RL 的混合策略，将不同思维模式（如逻辑推理、知识检索等）整合到统一框架中，增强模型的灵活性。
- 通用强化学习：以 RL 为主导，对模型进行全局优化，强化其在多样化任务中的通用性和鲁棒性。

Qwen3 使用

使用上和 Qwen2.5 没有太大区别，主要是增加了 Reasoning 模式的选择开关
这个开关也非常简便，大概有两种方式可以选择：
1. 在最新版 Transformer 库中 tokenizer.apply_chat_template 中可以设置 enable_thinking=True # Switch between thinking and non-thinking modes. Default is True.
2. 另外一种方式就是始终保持上面的开关常开，然后在 prompt 中添加 /no_think 来关闭推理模式。/no_think 可加到 system content 或 user content 结尾。
在一些下游任务上做了 SFT 微调测试，发现 Qwen3-1.7B 经过 SFT 之后和 Qwen2.5-3B-Instruct 经过 SFT 之后的效果差不多，这一点和官方对预训练的说法一致

Thoughts

在当今这个大模型结构同质化时代，预训练数据越多，模型的能力越强（同样效果下数据翻一倍模型尺寸可缩减一半！！），所以小厂没有这么多数据工程师，搞不到大量的高质量的数据，是很难在大模型上追赶大厂
到今天为止，Llama 系列官方模型都不能做到原生支持中文，和 Qwen 系列模型原生支持 119 种语言和方言相比，感觉非常小家子气，注定会被扫进历史的垃圾堆里
DeepSeek 对 Qwen 的影响很大，比如：
- 从 DeepSeek v3 之后，Qwen 系列模型预训练和最终版的命名从：Qwen2.5-7B/Qwen2.5-7B-Instruct 变成了 Qwen3-8B-Base/Qwen3-8B
- 广泛的使用了模型蒸馏，确实极大的提高了小尺寸模型的能力（之前 Qwen2.5-0.5B-Instruct 基本就是个答非所问的傻子，现如今的 Qwen3-0.6B 在不开推理模式的情况下也可以解决很多数学问题，非常强）
- 四段式后训练也和 DeepSeek 使用的后训练非常相似
合并推理和非推理模型的做法，是今后大模型的趋势，有不少模型都在朝着这个方向发展
官方表示 Qwen 系列后续会向着 Agent 方向发展，铺垫了很久的 MCP 可能会产生新的 Agent 应用变革

CUDA 学习笔记 01 —— CUDA 基础

Posted on 2025-03-24 Edited on 2025-05-11 In CUDA Valine:

0. 学习 CUDA 的目的

作为一个算法工程师，平时接触 HPC (High Performance Computing) 的机会并不多，那为什么还要学习 CUDA 呢？
学习 CUDA 的目的不是为了用 CUDA 做模型加速，而是从 CUDA 角度理解目前较新的大模型设计理念，这些高性能模型是如何从原理上做到又快又好的。
例如火出圈的 DeepSeek 系列模型，在模型设计角度做了较多创新，并开源了部分 CUDA 代码，对于不了解 CUDA 的工程师，很难 get 到算法设计的优雅之处。
反观某家大模型基座公司，曾开源某个模型结构，论文中一通自夸，分析理论计算量有多低。但很多人实测发现速度并没有很快，究其原因，实际上是这家公司还用的小模型时代的旧思维，即：一个模型理论计算量低，那就是快。
大模型时代不了解硬件，不尊重硬件，在算法创新上不太可能走的远。

1. Hello World

cuda 代码

#include <cuda_runtime_api.h>
#include <stdio.h>

// cuda 中 host 表示 cpu 端，device 表示 gpu 端
// __device__ 是设备函数的声明符号，表明该函数在 device 执行，且只能在 device
// 中调用
__device__ const char *device_hello_world(void) {
  return "GPU: Hello world!\n";
}

// __host__ 是主机函数的声明符号，表明该函数在 host 执行，且只能在 host 中调用
__host__ const char *host_hello_world(void) { return "CPU: Hello world!\n"; }

// __global__ 是核函数的声明符号，表明该函数在 device 执行，且只能在 host 中调用
__global__ void hello_world(void) {
  const char *str = device_hello_world();
  printf("%s", str);
}

int main(int argc, char **argv) {
  printf("%s", host_hello_world());
  // <<<grid_dim, block_dim>>> 是核函数的调用符号，表示启动 grid_dim 个 block，
  // 每个 block 有 block_dim 个线程
  hello_world<<<1, 10>>>();
  cudaDeviceReset();
  return 0;
}

cuda 的三个函数声明符号：
- __host__：主机函数，表示该函数在 CPU 上执行，且只能在 CPU 中调用
- __device__：设备函数，表示该函数在 GPU 上执行，且只能在 GPU 中调用
- __global__：核函数，表示该函数在 GPU 上执行，且只能在 CPU 中调用
其中 __global__ 声明的函数类型被称为 核函数，是 CUDA 中最重要的函数类型
- 核函数通过 <<<grid_dim, block_dim>>> 的方式调用，其中 <<<>>> 是 cuda 扩展关键字
- grid_dim 表示启动的 block 数量，block_dim 表示每个 block 中的线程数量
- grid_dim 和 block_dim 都是 dim3 类型的变量，表示三维数组，如果使用整形则模型 y 和 z 维度都为 1

编译

1	nvcc hello_world.cu -g -o hello_world

运行

./hello_world
# CPU: Hello world!
# GPU: Hello world!
# GPU: Hello world!
# GPU: Hello world!
# GPU: Hello world!
# GPU: Hello world!
# GPU: Hello world!
# GPU: Hello world!
# GPU: Hello world!
# GPU: Hello world!
# GPU: Hello world!

2. dimension 测试

#include <cuda_runtime.h>
#include <stdio.h>
__global__ void checkIndex(void) {
  printf("threadIdx:(%d,%d,%d) blockIdx:(%d,%d,%d) blockDim:(%d,%d,%d)\
  gridDim(%d,%d,%d)\n",
         threadIdx.x, threadIdx.y, threadIdx.z, blockIdx.x, blockIdx.y,
         blockIdx.z, blockDim.x, blockDim.y, blockDim.z, gridDim.x, gridDim.y,
         gridDim.z);
}
int main(int argc, char **argv) {
  int nElem = 6;  // number of elements
  dim3 block(3);  // block size
  int nBlock = (nElem + block.x - 1) / block.x; // number of blocks
  dim3 grid(nBlock);  // grid size
  printf("grid.x %d grid.y %d grid.z %d\n", grid.x, grid.y, grid.z);  
  printf("block.x %d block.y %d block.z %d\n", block.x, block.y, block.z);
  checkIndex<<<grid, block>>>();
  cudaDeviceReset();
  return 0;
}

执行结果：

./check_dimension
# grid.x 2 grid.y 1 grid.z 1
# block.x 3 block.y 1 block.z 1
# threadIdx:(0,0,0) blockIdx:(1,0,0) blockDim:(3,1,1)  gridDim(2,1,1)
# threadIdx:(1,0,0) blockIdx:(1,0,0) blockDim:(3,1,1)  gridDim(2,1,1)
# threadIdx:(2,0,0) blockIdx:(1,0,0) blockDim:(3,1,1)  gridDim(2,1,1)
# threadIdx:(0,0,0) blockIdx:(0,0,0) blockDim:(3,1,1)  gridDim(2,1,1)
# threadIdx:(1,0,0) blockIdx:(0,0,0) blockDim:(3,1,1)  gridDim(2,1,1)
# threadIdx:(2,0,0) blockIdx:(0,0,0) blockDim:(3,1,1)  gridDim(2,1,1)

3. CUDA 向量加法

代码

#include "freshman.h"
#include <cuda_runtime.h>
#include <stdio.h>

__host__ void sumArrays(float *a, float *b, float *res, const int size) {
  for (int i = 0; i < size; i++) {
    res[i] = a[i] + b[i];
  }
}
__global__ void sumArraysGPU(float *a, float *b, float *res, const int size) {
  int i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i < size) // 线程索引越界检查
    res[i] = a[i] + b[i];
}
int main(int argc, char **argv) {
  // set up device
  initDevice(0);

  // allocate host memory
  int nElem = 1 << 24;
  printf("Vector size:%d\n", nElem);
  int nByte = sizeof(float) * nElem;
  float *a_h = (float *)malloc(nByte);
  float *b_h = (float *)malloc(nByte);
  float *res_h = (float *)malloc(nByte);
  float *res_from_gpu_h = (float *)malloc(nByte);
  memset(res_h, 0, nByte);
  memset(res_from_gpu_h, 0, nByte);

  // allocate device memory
  float *a_d, *b_d, *res_d;
  CHECK(cudaMalloc((float **)&a_d, nByte));
  CHECK(cudaMalloc((float **)&b_d, nByte));
  CHECK(cudaMalloc((float **)&res_d, nByte));

  // randomly initialize the input data
  initialData(a_h, nElem);
  initialData(b_h, nElem);
  CHECK(cudaMemcpy(a_d, a_h, nByte, cudaMemcpyHostToDevice));
  CHECK(cudaMemcpy(b_d, b_h, nByte, cudaMemcpyHostToDevice));

  // set up execution configuration
  // 1. 计算最佳块大小
  int minGridSize, bestBlockSize;
  cudaOccupancyMaxPotentialBlockSize(&minGridSize, &bestBlockSize,
                                     (void *)sumArraysGPU,
                                     0, // 动态共享内存大小
                                     0  // 无块大小限制
  );

  // 2. 设置网格和块维度
  dim3 block(bestBlockSize);
  dim3 grid((nElem + bestBlockSize - 1) / bestBlockSize);

  // 3. 设备执行并统计耗时
  double iStart, iElaps;
  iStart = cpuSecond();
  sumArraysGPU<<<grid, block>>>(a_d, b_d, res_d, nElem);
  iElaps = cpuSecond() - iStart;
  CHECK(cudaGetLastError());
  CHECK(cudaDeviceSynchronize());
  CHECK(cudaMemcpy(res_from_gpu_h, res_d, nByte, cudaMemcpyDeviceToHost));
  printf("Execution configuration<<<%d,%d>>> Time elapsed %f sec\n", grid.x,
         block.x, iElaps);

  // 4. CPU执行并统计耗时
  iStart = cpuSecond();
  sumArrays(a_h, b_h, res_h, nElem);
  iElaps = cpuSecond() - iStart;
  printf("CPU Time elapsed %f sec\n", iElaps);

  // 5. 检查结果
  checkResult(res_h, res_from_gpu_h, nElem);

  // 6. 释放内存
  cudaFree(a_d);
  cudaFree(b_d);
  cudaFree(res_d);

  free(a_h);
  free(b_h);
  free(res_h);
  free(res_from_gpu_h);

  return 0;
}

cudaOccupancyMaxPotentialBlockSize 函数用于计算最佳块大小
max thread per block 是 cuda 中的一个限制，表示每个块中最多可以有多少个线程，一般为 1024，当超过这个限制时，CHECK(cudaGetLastError()); 会报错

运行结果

./sum_arrays
# Using device 0: NVIDIA GeForce RTX 4060 Laptop GPU
# Vector size:16777216
# Execution configuration<<<21846,768>>> Time elapsed 0.000030 sec
# CPU Time elapsed 0.076604 sec
# Check result success!

可以看出，cuda 的执行时间远远小于 cpu 的执行时间，相差了 2553 倍

总体流程

graph TD
    A[Host 端申请内存] --> B[Host 端初始化输入数据]
    B --> C[Device 端申请内存]
    C --> D[拷贝 Host 输入数据到 Device 端]
    D --> E[Device 端执行核函数]
    E --> F[拷贝 Device 端输出数据到 Host 端]
    F --> G[Host 端检查结果]
    G --> H[释放 Host 端和 Device 端全部内存]
    B --> J[Host 端执行普通函数] --> G

细节分析

1. `cuda` 内存分配是怎么做的？

1 2	float a_d; // 空指针 cudaMalloc((float *)&a_d, nByte); // 将指针的地址转成二级指针（指针的指针）传入内存分配函数

这里的二级指针应用很巧妙，由于 c++ 中的指针是值传递，所以如果是一级指针传入 cudaMalloc 函数时，指针 a_d 的值不会改变，因此只能将指针的地址转成二级指针传入内存分配函数

2. `cuda` 内存拷贝是怎么做的？

1	cudaMemcpy(a_d, a_h, nByte, cudaMemcpyHostToDevice);

cudaMemcpy 函数的四个参数分别是：
- dst：目标地址
- src：源地址
- size：拷贝的字节数
- kind：拷贝的类型，属于 cudaMemcpyKind 枚举类型
  - cudaMemcpyHostToHost = 0, /**< Host -> Host */
  - cudaMemcpyHostToDevice = 1, /**< Host -> Device */
  - cudaMemcpyDeviceToHost = 2, /**< Device -> Host */
  - cudaMemcpyDeviceToDevice = 3, /**< Device -> Device */
  - cudaMemcpyDefault = 4 /**< Direction of the transfer is inferred from the pointer values. Requires unified virtual addressing */

3. `cuda` 核函数如何解析线程索引？

1
2
3

int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < size) // 线程索引越界检查
  res[i] = a[i] + b[i];

blockIdx：表示当前块的索引
blockDim：表示当前块的维度（每个块中的线程数）
threadIdx：表示当前线程的索引
每个线程中计算两个标量的和
由于 gridDim * blockDim 可能大于 size，所以需要判断线程索引是否越界

4. 如何计算最佳块大小？

int minGridSize, bestBlockSize;
cudaOccupancyMaxPotentialBlockSize(&minGridSize, &bestBlockSize,
                                     (void *)sumArraysGPU,
                                     0, // 动态共享内存大小
                                     0  // 无块大小限制
);

cudaOccupancyMaxPotentialBlockSize 函数用于计算最佳块大小，五个参数分别是：
- minGridSize：最小网格大小变量地址
- bestBlockSize：最佳块大小变量地址
- kernel：核函数指针
- dynamicSMemSize：动态共享内存大小
- blockSizeLimit：块大小限制
函数名就是函数地址，可强转为 void * 函数指针（也可以写成：(void *)&sumArraysGPU）

4. CUDA 编程模型

线程块

线程块 block 是 CUDA 中的逻辑执行单元，是一个三维逻辑结构：
- block.x：表示块的 x 维度大小
- block.y：表示块的 y 维度大小
- block.z：表示块的 z 维度大小
- 其中 block.x 是最内层的循环，block.y 是第二层循环，block.z 是最外层的循环
- 用三维数组可以表示为：tread[z][y][x]，即 tid = threadIdx.x + threadIdx.y * blockDim.x + threadIdx.z * blockDim.x * blockDim.y

线程束

线程束 warp 是 CUDA 基本调度执行单元，一个 warp 由 32 个线程组成
- 一个 warp 中的线程在一个时钟周期内执行同一条指令（单指令多线程，SIMT）
- 一个 warp 中的线程可以共享指令指针和执行资源（如寄存器、缓存等）
- Warp 调度器（warp scheduler）负责将 warp 分配到物理执行单元上执行
线程块会被划分为多个 warp， $WarpsPerBlock=ceil(\frac{ThreadsPerBlock}{32})$

CUDA core

CUDA core 是 CUDA 物理执行单元，负责实际的计算任务
一个 CUDA core 一个时钟周期只能计算一个线程的指令

StreamMultiprocessor

StreamMultiprocessor 流式多处理器（简称 SM），负责执行 CUDA 线程块中的并行计算任务
每个 GPU 包含多个 SM，每个 SM 包含多个 CUDA core，例如：RTX 4060 有 24 个 SM，每个 SM 有 128 个 CUDA core

5. Reduce

规约（Reduce）是 CUDA 编程中常见的操作，主要用于将多个数据元素规约为一个数据元素
规约操作通常是一个二元操作，例如：sum、mul、max、min 等，简单的规约可以合并成强大的算子，甚至可以说规约算子是神经网络的基础

规约求和

// CPU 规约求和
int recursiveReduce(int *data, int const size) {
  // terminate check
  if (size == 1)
    return data[0];
  // renew the stride
  int const stride = size / 2;
  if (size % 2 == 1) {
    for (int i = 0; i < stride; i++) {
      data[i] += data[i + stride];
    }
    data[0] += data[size - 1];
  } else {
    for (int i = 0; i < stride; i++) {
      data[i] += data[i + stride];
    }
  }
  // call
  return recursiveReduce(data, stride);
}
// GPU 规约相邻求和
__global__ void reduceNeighbored(int *g_idata, int *g_odata, unsigned int n) {
  // set thread ID
  unsigned int tid = threadIdx.x;
  // boundary check
  if (tid >= n)
    return;
  // convert global data pointer to local point of this block
  int *idata = g_idata + blockIdx.x * blockDim.x;
  // in-place reduction in global memory
  for (int stride = 1; stride < blockDim.x; stride *= 2) {
    if ((tid % (2 * stride)) == 0) {
      idata[tid] += idata[tid + stride];
    }
    // synchronize within block
    __syncthreads();
  }
  // write result for this block to global mem
  if (tid == 0)
    g_odata[blockIdx.x] = idata[0];
}
// GPU 规约相邻求和（简化版）
__global__ void reduceNeighboredLess(int *g_idata, int *g_odata,
                                     unsigned int n) {
  unsigned int tid = threadIdx.x;
  unsigned idx = blockIdx.x * blockDim.x + threadIdx.x;
  // convert global data pointer to the local point of this block
  int *idata = g_idata + blockIdx.x * blockDim.x;
  if (idx > n)
    return;
  // in-place reduction in global memory
  for (int stride = 1; stride < blockDim.x; stride *= 2) {
    // convert tid into local array index
    int index = 2 * stride * tid;
    if (index < blockDim.x) {
      idata[index] += idata[index + stride];
    }
    __syncthreads();
  }
  // write result for this block to global men
  if (tid == 0)
    g_odata[blockIdx.x] = idata[0];
}
// GPU 规约交错求和，主要是 stride 的计算方式不同
__global__ void reduceInterleaved(int *g_idata, int *g_odata, unsigned int n) {
  unsigned int tid = threadIdx.x;
  unsigned idx = blockIdx.x * blockDim.x + threadIdx.x;
  // convert global data pointer to the local point of this block
  int *idata = g_idata + blockIdx.x * blockDim.x;
  if (idx >= n)
    return;
  // in-place reduction in global memory
  for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {

    if (tid < stride) {
      idata[tid] += idata[tid + stride];
    }
    __syncthreads();
  }
  // write result for this block to global men
  if (tid == 0)
    g_odata[blockIdx.x] = idata[0];
}

效率分析

本代码中使用了三种核函数实现方式做同一个规约操作，分别是：
- reduceNeighbored：相邻线程规约
- reduceNeighboredLess：相邻线程规约（简化版）
- reduceInterleaved：交错线程规约
三者效率从高到低依次是：
- reduceInterleaved > reduceNeighboredLess > reduceNeighbored
三者的示意图分别如下：
- reduceNeighbored：相邻线程规约的实现
- reduceNeighboredLess：相邻线程规约的简化版实现（注意圆圈中的符号已和上图不一致）
- reduceInterleaved：交错线程规约的实现
三者效率差异主要来自于 线程分支分化，后续会详细介绍

6. 循环展开

循环展开（Loop Unrolling）是 CUDA 中常用的优化手段，主要用于减少循环控制开销和提高指令级并行度
简单说就是一个线程不再只计算一个数据，而是计算多个数据，而且是直接在代码中展开，而不是在编译器中展开
可以简单理解成：启动线程是需要花时间的，启动一个线程只算一个数据，太浪费了，所以我们可以让一个线程计算多个数据，这样就能减少启动线程的时间开销，所以就省时间了

代码

// 数据总长是 8 * blockDim.x * gridDim.x
// 线程数是 blockDim.x * gridDim.x
// 每个线程计算 8 个数据
__global__ void reduceUnroll8(int *g_idata, int *g_odata, unsigned int n) {
  unsigned int tid = threadIdx.x;
  unsigned int idx = blockDim.x * blockIdx.x * 8 + threadIdx.x;
  if (tid >= n)
    return;
  int *idata = g_idata + blockIdx.x * blockDim.x * 8;
  // 循环展开，每个线程计算 8 个数据
  // 直接把 8 * blockDim.x * gridDim.x 的数据总长
  // 聚合到了 blockDim.x * gridDim.x 的线程数上
  if (idx + 7 * blockDim.x < n) {
    g_idata[idx] += g_idata[idx + blockDim.x];
    g_idata[idx] += g_idata[idx + blockDim.x * 2];
    g_idata[idx] += g_idata[idx + blockDim.x * 3];
    g_idata[idx] += g_idata[idx + blockDim.x * 4];
    g_idata[idx] += g_idata[idx + blockDim.x * 5];
    g_idata[idx] += g_idata[idx + blockDim.x * 6];
    g_idata[idx] += g_idata[idx + blockDim.x * 7];
  }
  // 这里需要同步，也就是线程阻塞直到所有线程都执行完
  __syncthreads();
  // 然后就是一个最简单的规约操作了，和上面一样
  for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
    if (tid < stride) {
      idata[tid] += idata[tid + stride];
    }
    __syncthreads();
  }
  // write result for this block to global mem
  if (tid == 0)
    g_odata[blockIdx.x] = idata[0];
}

7. 核函数递归调用

和 CPU 一样，CUDA 也支持核函数递归调用，调用方式和普通递归函数一样
需要注意的是在编译的时候需要加上 -rdc=true 选项

代码

__global__ void nesthelloworld(int iSize, int iDepth) {
  unsigned int tid = threadIdx.x;
  printf("depth : %d blockIdx: %d,threadIdx: %d\n", iDepth, blockIdx.x,
         threadIdx.x);
  if (iSize == 1)
    return;
  int nthread = (iSize >> 1);
  if (tid == 0 && nthread > 0) {
    // 递归调用核函数
    nesthelloworld<<<1, nthread>>>(nthread, ++iDepth);
    printf("-----------> nested execution depth: %d\n", iDepth);
  }
}

8. 固定内存

Pinned Memory 是 CUDA 中的一种特殊内存类型（不是显存，是内存），主要用于提高数据传输效率
普通内存是分页管理，分页管理存在两个问题：
1. 一页内存逻辑上连续，但物理上不连续
2. 操作系统可能会将内存页交换到磁盘上，导致数据不在物理内存中
Pinned Memory 就是解决了这两个问题，分配了一块连续物理地址且固定的主机内存（host 内存），方便整块拷贝数据到显存（DMA）

代码

#include "freshman.h"
#include <cuda_runtime.h>

int main(int argc, char **argv) {
  int dev = 0;
  cudaSetDevice(dev);

  int nElem = 1 << 14;
  int nByte = sizeof(float) * nElem;
  float *a_h, *b_h, *res_h, *res_from_gpu_h;
  // 注意这里的 cudaMallocHost 和 cudaMalloc 是不同的
  // 前者申请的是 host 固定内存，后者申请的是 device 显存
  // cudaMallocHost 是 malloc 的一个平替
  CHECK(cudaMallocHost((float **)&a_h, nByte));
  CHECK(cudaMallocHost((float **)&b_h, nByte));
  CHECK(cudaMallocHost((float **)&res_h, nByte));
  CHECK(cudaMallocHost((float **)&res_from_gpu_h, nByte));

  // 初始化数据
  memset(res_h, 0, nByte);
  memset(res_from_gpu_h, 0, nByte);
  initialData(a_h, nElem);
  initialData(b_h, nElem);

  // 申请设备显存
  float *a_d, *b_d, *res_d;
  CHECK(cudaMalloc((float **)&a_d, nByte));
  CHECK(cudaMalloc((float **)&b_d, nByte));
  CHECK(cudaMalloc((float **)&res_d, nByte));

  // 拷贝数据到设备显存
  CHECK(cudaMemcpy(a_d, a_h, nByte, cudaMemcpyHostToDevice));
  CHECK(cudaMemcpy(b_d, b_h, nByte, cudaMemcpyHostToDevice));

  // 跑核函数
  dim3 block(1024);
  dim3 grid(nElem / block.x);
  sumArraysGPU<<<grid, block>>>(a_d, b_d, res_d);
  printf("Execution configuration<<<%d,%d>>>\n", grid.x, block.x);

  // 结果拷贝回主机并检查两个设备计算结果是否一致
  CHECK(cudaMemcpy(res_from_gpu_h, res_d, nByte, cudaMemcpyDeviceToHost));
  sumArrays(a_h, b_h, res_h, nElem);
  checkResult(res_h, res_from_gpu_h, nElem);

  // 释放内存
  // 注意这里的 cudaFreeHost 和 cudaFree 是不同的
  // 前者释放的是 host 固定内存，后者释放的是 device 显存
  cudaFree(a_d);
  cudaFree(b_d);
  cudaFree(res_d);
  cudaFreeHost(a_h);
  cudaFreeHost(b_h);
  cudaFreeHost(res_h);
  cudaFreeHost(res_from_gpu_h);

  return 0;
}

关键点

主存 普通内存 的分配和释放函数是 malloc 和 free
主存 固定内存 的分配和释放函数是 cudaMallocHost 和 cudaFreeHost
显存的分配和释放函数是 cudaMalloc 和 cudaFree

9. 零拷贝内存和统一虚拟地址

Zero-Copy Memory 是 CUDA 中一种允许 GPU 直接访问主机内存的技术，避免了显式的数据拷贝操作（不需要用 cudaMemcpy 函数）
实际上，Zero-Copy Memory 在很多时候并不快，因为 GPU 访问主机内存的速度远远低于访问显存的速度，因此，Zero-Copy Memory 只适用于一些特殊的场景，例如：
- 主机内存中的数据只需要被 GPU 使用一次
- 数据量太大，显存放不下
- 调试用途
Zero-Copy Memory 的实现方式是将主机内存映射到 GPU 的地址空间中，GPU 通过访问这个地址空间来访问主机内存，实际上走的是 PCIe 总线
由于不需要先完成所有数据的拷贝再开始执行核函数，因此 Zero-Copy Memory 使用异步拷贝的方式来实现，可将部分拷贝数据的时间和核函数执行的时间重叠，但并不多
Unified Virtual Addressing (UVA) 是 CUDA 中的一种内存管理机制，允许 CPU 和 GPU 共享同一虚拟地址空间

代码

float *a_host, *b_host, *res_d;
// 申请主机固定内存，添加特殊 flag cudaHostAllocMapped
CHECK(cudaMallocHost((float **)&a_host, nByte, cudaHostAllocMapped));
CHECK(cudaMallocHost((float **)&b_host, nByte, cudaHostAllocMapped));

// a_host 和 b_host 是可直接作为核函数的输入参数
// 也可以转成 Device 地址空间，如下：
float *a_dev, *b_dev;
// 映射主机内存到设备地址空间
CHECK(cudaHostGetDevicePointer((void **)&a_dev, (void *)a_host, 0));
CHECK(cudaHostGetDevicePointer((void **)&b_dev, (void *)b_host, 0));

// 用 a_dev 和 b_dev 作为核函数的输入参数计算

和 Zero-Copy Memory 一样，UVA 甚至不需要将 a_host 转成 a_dev，直接用 a_host 就可以调用核函数

10. Aos 和 SoA

CUDA 不仅支持最简单的原生数据类型，还支持自定义数据类型（struct），例如 Aos 和 SoA 等
Aos（Array of Structures）和 SoA（Structure of Arrays）是两种不同的数据存储方式，这两种方式由于变量的排布方式不同，导致了访问内存的效率差异

Aos

Aos 是将多个结构体存储在一个数组中，每个结构体的成员变量是连续存储的
例如：

struct AoSStruct {
  float a;
  float b;
};
__global__ void sumArraysGPU(float *a, float *b, struct naiveStruct *res, int n) {
  int i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i < n)
    res[i].a = a[i] + b[i];
}

这里的 res 是一个 结构体数组，a 和 b 是结构体的成员变量，每个变量都是一个 float 类型的标量
每个结构体的成员变量是连续存储的，即：a1 b1 a2 b2 a3 b3 ...

SoA

SoA 是将多个结构体的成员变量存储在一个数组中，每个成员变量是连续存储的
例如：

struct SoAStruct {
  float a[SIZE];
  float b[SIZE];
};
__global__ void sumArraysGPU(float *a, float *b, struct SoAStruct *res, int n) {
  int i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i < n)
    res->a[i] = a[i] + b[i];
}

这里的 res 是一个 结构体，a 和 b 是结构体的成员变量，每个变量都是一个 float 类型的数组
每个成员变量的数组是连续存储的，即：a1 a2 a3 ... b1 b2 b3 ...

11. 行主序和列主序

行主序（Row Major Order）和 列主序（Column Major Order）是两种不同的数组存储方式
行主序 是将数组的每一行存储在连续的内存中，列主序 是将数组的每一列存储在连续的内存中
例如：

int a[3][4] = {
    {1, 2, 3, 4},
    {5, 6, 7, 8},
    {9, 10, 11, 12}};

行主序 存储方式是：1 2 3 4 5 6 7 8 9 10 11 12
列主序 存储方式是：1 5 9 2 6 10 3 7 11 4 8 12
默认情况下，C / C++ / CUDA 语言是 行主序 存储方式
在行主序存储下，如果按行序访问数组元素，访问效率会更高，因为连续的内存访问会提高缓存命中率，反之如果按列序访问数组元素，访问效率会更低

速度对比

WIP

Transformers without Normalization

Posted on 2025-03-17 Edited on 2025-03-21 In Transformer Valine:

URL

paper: https://arxiv.org/pdf/2503.10622
code: https://github.com/jiachenzhu/DyT

TL;DR

这是由恺明和杨立昆提出的一篇关于 transformer 算子优化的论文，主要观点是去掉 transformer 结构中的 normalization 层，改成 tanh 层
改用 tanh 算子的 transformer 模型，在大多数任务上可达到使用归一化层的模型相同的性能，甚至更好

Algorithm

简单来说，这篇论文的核心思想是将 transformer 中的 normalization 层（可以是 LayerNorm 或 RMSNorm）替换成 dynamic tanh 层（简称 DyT）
normalization 计算公式：

\text{normalization}(x) = \gamma \times \frac{x - \mu}{\sqrt{\sigma^2+\epsilon}} + \beta

其中 $\mu$ 和 $\sigma$ 分别是 mean 和 std， $\gamma$ 和 $\beta$ 是 scale 和 shift 参数

DyT 计算公式：

\text{DyT}(x) = \gamma \times \tanh(\alpha x) + \beta

其中 $\alpha$ 是个可学习参数， $\gamma$ 和 $\beta$ 是 scale 和 shift 参数（和 normalization 一样）

DyT 实现伪代码：

# input x has the shape of [B, T, C]
# B: batch size, T: tokens, C: dimension
class DyT(Module):
    def __init__(self, C, init_α):
        super().__init__()
        self.α = Parameter(ones(1) * init_α)
        self.γ = Parameter(ones(C))
        self.β = Parameter(zeros(C))
    def forward(self, x):
        x = tanh(self.alpha * x)
        return self.γ * x + self.β

$\alpha$ 默认初始化值为 0.5

Results

作者在多个领域的知名模型上都对比了修改前后训练精度，DyT 的性能和 normalization 的性能基本一致，打的有来有回
作者还对比了 DyT 和 normalization 的训练/推理速度，DyT 的训练/推理速度要快很多
作者同时做 tanh 和 $\alpha$ 做了消融实验，发现 tanh 和 $\alpha$ 都是必要的

Thoughts

属于是恺明和立昆的梦幻联动了…，这种对最火的结构的优化，非大佬不能为也，想象下如果这篇论文是大学实验室发表的，大家第一反应恐怕是：Who think you are? 😂
之前算是稍微接触过硬件，DyT 这种 element-wise op 比 normalization 这种 reduce op 一定快多了，想怎么 tiling 都行…

2025.02 DeepSeek 开源周第三弹 —— DeepGEMM

Posted on 2025-03-10 Edited on 2025-03-21 In LLM Valine:

WIP

2025.02 DeepSeek 开源周第二弹 —— DeepEP

Posted on 2025-03-06 Edited on 2025-03-21 In LLM Valine:

URL

code: https://github.com/deepseek-ai/DeepEP

TL;DR

目前大模型常用的三种并行：
- DP: Data Parallelism，数据并行，将数据分成多份，每个 GPU 处理一份数据
- PP: Pipeline Parallelism，管道并行，将模型分成多个阶段（连续一层或多层为一个阶段），每个 GPU 处理一个阶段
- TP: Tensor Parallelism，张量并行，将模型分成多份（通常是一层/一个算子/一个张量分成多份，主要解决超长序列引起的超大张量问题），每个 GPU 处理一部分张量
对于 MoE 结构，EP (Expert Parallelism) 是一种新的并行策略，将 MoE 中的 Expert 分配到不同的 GPU 上
DeepEP 是一个用 cuda 实现的 MoE 模型的并行库，重点在于对 All-to-All 通信的优化

背景知识

节点内通信

通俗讲：一台服务器被称为一个节点，一个节点上的多个 GPU 之间的通信被称为节点内通信
通信协议：
- PCIe：比较通用的通信协议，目前最新的 PCIe 6.0 的带宽为 32GB/s 的双向带宽
- NVLink：NVIDIA 自家的通信协议，目前最新的 NVLink 5.0 的带宽约为 800GB/s 的双向带宽

节点间通信

多台服务器组成一个集群，集群中的服务器之间的通信被称为节点间通信
通信协议：
- InfiniBand：HPC 领域常用的通信协议，目前最新的 InfiniBand NDR 可达 400Gbps 的双向带宽
- Ethernet：通用的通信协议，速度低于 InfiniBand
通信技术：
- RDMA：Remote Direct Memory Access，远程直接内存访问，RDMA 通过 DMA 直接访问远程内存，减少了 CPU 的参与，提高了通信效率

通信视角下 MoE 结构的特殊性

MoE 结构的本质是超大规模参数量 + 小规模激活参数量（实际计算量）来让模型更强大同时推理效率高
由于 MoE 结构在实际推理过程中，每个 token 激活的专家 id 是无法提前预测的，而是一个纯 runtime 的行为
对此，为了降低 EP 通信压力，MoE 结构通常会限制每个 token 实际激活的节点数量。例如，DeepSeek V3 有 1 个共享专家和 256 个路由专家，每个 token 会激活 1 个共享专家和 8 个路由专家，但同时限制最多只能激活 4 个节点，假如得分最高的 8 个路由专家来自超过 4 个节点，那么会牺牲部分高分专家，在节点数不超过 4 的情况下，用贪心算法选择得分最高的 4 个专家
虽然 DP / TP / PP 都存在通信问题，但都是可提前规划好的通信数据量和通信模式，只要调度得当，即可重叠计算和通信耗时，而 MoE 结构的通信是无法提前规划的，因此 MoE 结构的通信是最难优化的

关键特性和能力

DeepEP 的关键特性和能力包括：

高吞吐量节点内通信：使用 NVLink 优化节点内所有到所有通信的内核，实现高达 155 GB/s 的带宽。
高吞吐量节点间通信：使用 RDMA 实现高效的跨节点所有到所有通信，在不同的 EP 配置中保持大约 45 GB/s 的带宽。
低延迟内核：专用推理解码内核，分发操作延迟低至 163 微秒，组合操作延迟低至 318 微秒。
FP8 支持：原生支持低精度操作，包括 FP8 分发，与大型模型中量化趋势一致。
灵活的 GPU 资源控制：可配置的 SM 使用，用于计算 - 通信重叠，允许精细调整性能优化。
自适应路由支持：在低延迟内核中支持自适应路由，使复杂拓扑中的网络利用更高效。

技术实现

DeepEP 是用 C++ 和 CUDA 组件实现的，并带有 Python 接口。实现包括几个关键组件：

缓冲管理：核心 Buffer 类管理 NVLink 和 RDMA 的通信缓冲区，处理内存分配和同步。
通信内核：
- 训练和推理预填充的高吞吐量内核
- 推理解码的低延迟内核
- 支持节点内（NVLink）和节点间（RDMA）通信
事件管理：EventOverlap 类提供 CUDA 事件处理和计算 - 通信重叠的工具。
分发和组合操作：
- dispatch：将令牌特征发送到跨 GPU 的对应专家
- combine：从专家收集处理后的特征并返回到原始位置

Thoughts

大模型，尤其是基座大模型，拼的是基建
大模型时代不会再出现小模型时代经常出现的 理论计算量低但实际很慢的算法 了，GPU 上快才是真的快，不光要考虑计算，存储 / 通信也同时需要认真考虑
软硬件 co-design 是未来趋势，People who're serious about software should make their own hardware. 这句名言的含金量还在上升

2025.02 DeepSeek 开源周第一弹 —— FlashMLA

Posted on 2025-03-06 Edited on 2025-03-21 In LLM Valine:

URL

code: https://github.com/deepseek-ai/FlashMLA

TL;DR

FlashMLA 是针对 DeepSeek 提出的 MLA (Multi-head Latent Attention) 模块，在 Nvidia Hopper 架构 GPU 上的加速算法，尤其对边长序列推理做了优化
FlashMLA 和 MLA 的关系大概可以类比到 FlashAttention 和 Attention 的关系

关键特性和能力

仅对 Hopper (sm_90) 架构做优化：充分挖掘了硬件 (sm) 的计算能力和 Memory Hierachy 的存储 / IO 能力
支持可变长度序列：和现实世界中推理场景更贴合
分页 KV cache：使用 block size = 64 的分页存储（这里的 64 的含义是：一个 block 存储某一层某个头的连续 64 个 token 对应的 kv cache）
高性能带宽和算力：在 H800 SXM5 设备上，在内存带宽受限配置环境下，内存带宽可达 3000 GB/s，在算力受限配置下，算力可达 580 TFLOPS
支持多种精度：支持 BF16 和 FP16

技术实现

基于 Nvidia cutlass 库 + cuda 实现主要计算
用干净的 python API 包装，方便集成到 PyTorch-base 框架中
用元数据管理的方式支持变长序列

Thoughts

做大模型不懂 cuda 是不行了，cutlass 要开始学起来了…
FlashMLA 只在 sm_90 上实现了，其他显卡编译不通，~~DeepSeek 只打高端局~~

Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention

Posted on 2025-02-20 Edited on 2025-03-21 In LLM Valine:

URL

paper: https://arxiv.org/pdf/2502.11089

TL;DR

本文提出了一种新的稀疏注意力机制，称为 Native Sparse Attention，该机制在硬件上对齐，并且可以直接训练，而无需额外的稀疏化技术
目前没有开源代码，只能通过论文中的公式和描述来实现

Algorithm

总体流程

本质上是用不同的 pattern 组合来替代 full attention，以减少计算量

代码实现

由 DeepSeek-R1 根据论文中的公式和描述实现

import torch
import torch.nn as nn
import torch.nn.functional as F
class NSAModule(nn.Module):
    def __init__(
        self,
        d_model,
        n_heads,
        compress_block=32,
        select_block=64,
        select_topk=16,
        window_size=512,
        dropout=0.1,
    ):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads
        # 参数设置
        self.compress_block = compress_block
        self.select_block = select_block
        self.select_topk = select_topk
        self.window_size = window_size
        # 压缩用MLP
        self.compress_mlp = nn.Sequential(
            nn.Linear(self.head_dim * compress_block, 256),
            nn.GELU(),
            nn.Linear(256, self.head_dim),
        )
        # 门控机制
        self.gate_mlp = nn.Sequential(nn.Linear(d_model, 3 * n_heads), nn.Sigmoid())
        # 投影层
        self.q_proj = nn.Linear(d_model, d_model)
        self.k_proj = nn.Linear(d_model, d_model)
        self.v_proj = nn.Linear(d_model, d_model)
        self.out_proj = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)
    def _compress_tokens(self, k, v):
        """压缩KV序列到块级别"""
        # 调整输入维度 (batch, seq_len, n_heads, head_dim)
        b, t, nh, hd = k.shape  # 修改为四维解包
        block_size = self.compress_block
        num_blocks = (t + block_size - 1) // block_size
        pad_len = num_blocks * block_size - t
        # 填充并分块
        k = F.pad(k, (0, 0, 0, 0, 0, pad_len))  # 添加头部维度的填充
        v = F.pad(v, (0, 0, 0, 0, 0, pad_len))
        # 调整维度: [batch, num_blocks, block_size, n_heads, head_dim]
        k_blocks = k.view(b, num_blocks, block_size, nh, hd)
        v_blocks = v.view(b, num_blocks, block_size, nh, hd)
        # 压缩处理 (保持头部分离)
        k_compressed = self.compress_mlp(
            k_blocks.permute(0, 1, 3, 2, 4).flatten(3)
        )  # [b, num_blocks, nh, hd]
        v_compressed = self.compress_mlp(v_blocks.permute(0, 1, 3, 2, 4).flatten(3))
        return k_compressed, v_compressed
    def _select_blocks(self, q, k_compressed, v_compressed):
        """基于注意力分数选择关键块"""
        # 计算压缩注意力分数
        scores = torch.einsum("bthd,bkhd->bthk", q, k_compressed) / (
            self.head_dim ** 0.5
        )
        probs = F.softmax(scores, dim=-1)
        # 选择topk块
        topk_scores, topk_indices = torch.topk(
            probs.mean(dim=2), self.select_topk, dim=-1
        )
        # 收集选中的块
        k_selected = torch.gather(
            k_compressed,
            1,
            topk_indices.unsqueeze(-1).expand(-1, -1, -1, self.head_dim),
        )
        v_selected = torch.gather(
            v_compressed,
            1,
            topk_indices.unsqueeze(-1).expand(-1, -1, -1, self.head_dim),
        )
        return k_selected, v_selected
    def forward(self, x, attn_mask=None):
        b, t, d = x.shape
        # 投影QKV并保持四维结构
        q = self.q_proj(x).view(b, t, self.n_heads, self.head_dim)
        k = self.k_proj(x).view(b, t, self.n_heads, self.head_dim)
        v = self.v_proj(x).view(b, t, self.n_heads, self.head_dim)
        # 压缩路径
        k_compressed, v_compressed = self._compress_tokens(k, v)
        # 选择路径
        k_selected, v_selected = self._select_blocks(q, k_compressed, v_compressed)
        # 滑动窗口
        k_window = k[:, max(0, t - self.window_size) :]
        v_window = v[:, max(0, t - self.window_size) :]
        # 门控权重
        gate = self.gate_mlp(x).view(b, t, 3, self.n_heads).permute(2, 0, 3, 1, 2)
        # 三路注意力计算
        attn_outputs = []
        for branch_k, branch_v in [
            (k_compressed, v_compressed),
            (k_selected, v_selected),
            (k_window, v_window),
        ]:
            scores = torch.einsum("bthd,bkhd->bthk", q, branch_k) / (
                self.head_dim ** 0.5
            )
            if attn_mask is not None:
                scores = scores.masked_fill(attn_mask == 0, -1e9)
            probs = F.softmax(scores, dim=-1)
            probs = self.dropout(probs)
            output = torch.einsum("bthk,bkhd->bthd", probs, branch_v)
            attn_outputs.append(output)
        # 门控融合
        weighted = sum(g * o for g, o in zip(gate, attn_outputs))
        output = weighted.contiguous().view(b, t, d)
        return self.out_proj(output)
def test_nsa_attention_shapes():
    # 测试参数
    batch_size = 2
    seq_len = 128
    d_model = 256
    n_heads = 8
    nsa_attn = NSAModule(d_model, n_heads)
    x = torch.randn(batch_size, seq_len, d_model)
    print("输入形状:", x.shape)
    # 打印中间形状
    q = nsa_attn.q_proj(x).view(batch_size, seq_len, n_heads, -1)
    k = nsa_attn.k_proj(x).view(batch_size, seq_len, n_heads, -1)
    v = nsa_attn.v_proj(x).view(batch_size, seq_len, n_heads, -1)
    print("\nQ/K/V形状:", q.shape)  # 应该输出 [2, 128, 8, 32]
    k_comp, v_comp = nsa_attn._compress_tokens(k, v)
    print("压缩后KV形状:", k_comp.shape)  # [2, 4, 8, 32]
if __name__ == "__main__":
    test_nsa_attention_shapes()

Thoughts

速度快并不惊讶，但效果竟然比 Full Attention 还好，如果实验比较 solid，那么 NSA 确实比较有前途
由于没有开源代码，所以只能通过论文中的公式和描述来实现，这对于一些实验复现和工程应用来说是一个挑战

DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning

Posted on 2025-02-18 Edited on 2025-04-29 In LLM Valine:

URL

paper: https://arxiv.org/pdf/2501.12948
code: https://github.com/deepseek-ai/DeepSeek-R1

TL;DR

本文提出了三个模型和其对应的训练方法，目的是提高大模型的 CoT 推理能力，三个模型分别是：
- DeepSeek-R1-Zero：由 DeepSeek V3 base 直接通过 RL 训练得到（DeepSeek V3 base 是 DeepSeek V3 只经过 pretrain 的版本，即 DeepSeek V3 = DeepSeek V3 base + post-train ）
- DeepSeek-R1：由 DeepSeek V3 base 在合成的高质量 CoT 数据集上 SFT + RL 训练得到，效果优于 DeepSeek-R1-Zero
- Distill model：用相同的合成高质量 CoT 数据训练的开源模型（参数量小且非 MoE）
其中训练方法是本文讲述的重点，DeepSeek R1 中的 R 表示 Reasoning（推导）

Algorithm

总体流程

Benchmark

DeepSeek-R1-Zero

DeepSeek-R1-Zero 是用 DeepSeek v3 base (pre-train) 直接通过 RL 训练得到的，不经过 SFT
用了以下的 Template 让模型生成 CoT 推理结果：
奖励建模分成两个部分：
- Accuracy rewards: 用于评估模型生成的 CoT 推理结果的准确性
- Format rewards: 用于评估模型生成的 CoT 推理结果的格式是否符合上述 prompt 要求的 CoT 的格式
且 RL 过程不使用 NN reward model，而是使用了 Rule base reward
RL 过程中使用了 GRPO 算法
DeepSeek-R1-Zero 仅仅用上述简单的 CoT 性能提升流程，即可大幅提高 benchmark 上的指标（相比于 DeepSeek v3），但仍然存在一些问题，例如可读性差和语言混合

DeepSeek-R1

DeepSeek-R1 用了更复杂的数据处理流程和训练流程，主要分成四个阶段（流程图上的四个 stage）：
- 冷启动：用数千个来自于 DeepSeek-R1-Zero 输出且经过人工处理的 CoT 数据，对 DeepSeek V3 base 进行 SFT
- 面向推导的强化学习：用编码、数学、科学和逻辑推理等推理密集型任务数据，对冷启动微调后的模型进行 RL 训练，同时引入了语言一致性和可读性的奖励
- 拒绝采样和监督微调：用上一个 stage 得到的模型生成更广泛领域的推导 CoT 数据 600K 条，用 DeepSeek V3 模型推理部分非推导数据（例如：写作、翻译等）得到潜在的思维链数据 200K 条，用这 800K 条数据对 DeepSeek V3 base 进行 SFT
- 所有场景的强化学习：用了和 DeepSeek V3 中使用的 RL 相同的 pipeline，数据和 DeepSeek-R1-Zero 中的 RL 数据相同，对上一个 stage 得到的模型进行 RL 训练
综上所述，DeepSeek-R1 最重要的是 800K 包含 CoT 的 SFT 数据

Distill model

由于 DeepSeek R1 使用了 DeepSeek V3 相同的模型结构，因此包含 671B 参数，每个 token 激活 34B 参数，因此这个模型非常大，不适合部署
为了解决这个问题，本文提出了 Distill model，用开源的较小参数量的稠密模型（非 MoE 模型）在 800K 包含 CoT 的 SFT 数据 进行 SFT 微调
主要蒸馏了 Qwen 和 llama 模型，且完全不做 RL，即可大幅提高原模型在 Benchmark 上的指标

Thoughts

需要重点关注本论文提到的 CoT 数据合成方式，这是 DeepSeek R1 的核心
能从本文中可以看出 数据重要性远大于模型 和 Post-train 只包含一次 SFT 和一次 RL，比如论文中 stage 1 模型会大胆抛弃 DeepSeek R1 Zero 模型，而是退回到了 DeepSeek V3 base 模型，只用了 Zero 生成数据；stage 3 模型会大胆抛弃之前 stage 的模型，而是退回到了 DeepSeek V3 base 模型，只用了 stage 2 模型生成数据
合成数据是最重要的，合成高质量的数据是 DeepSeek R1 成功的关键，甚至可以说如果 DeepSeek 开源了 80K 条合成数据，那么 CoT 领域的研究将会有一个新的起点

Reference

https://zhuanlan.zhihu.com/p/22103896675

URL

TL;DR

Algorithm

DDPM 逆向过程公式

DDIM 逆向过程公式

DDIM 公式拆解：

1. 预测原输入

2. 计算调整方向

3. 更新隐变量

DDPM 和 DDIM 使用上的区别

二者对比分析

Thoughts

URL

TL;DR

URL

TL;DR

Qwen3 系列模型

Qwen3 模型的主要特性

预训练

后训练

Qwen3 使用

Thoughts

0. 学习 CUDA 的目的

1. Hello World

cuda 代码

编译

运行

2. dimension 测试

3. CUDA 向量加法

代码

运行结果

总体流程

细节分析

1. cuda 内存分配是怎么做的？

2. cuda 内存拷贝是怎么做的？

3. cuda 核函数如何解析线程索引？

4. 如何计算最佳块大小？

4. CUDA 编程模型

线程块

线程束

CUDA core

StreamMultiprocessor

5. Reduce

规约求和

效率分析

6. 循环展开

代码

7. 核函数递归调用

代码

8. 固定内存

代码

关键点

9. 零拷贝内存 和 统一虚拟地址

代码

10. Aos 和 SoA

Aos

SoA

11. 行主序和列主序

速度对比

URL

TL;DR

Algorithm

Results

Thoughts

WIP

URL

TL;DR

背景知识

节点内通信

节点间通信

通信视角下 MoE 结构的特殊性

关键特性和能力

技术实现

Thoughts

URL

TL;DR

关键特性和能力

技术实现

Thoughts

URL

`DDIM` 公式拆解：

1. `cuda` 内存分配是怎么做的？

2. `cuda` 内存拷贝是怎么做的？

3. `cuda` 核函数如何解析线程索引？

9. 零拷贝内存和统一虚拟地址