flash-attn 实现 prefill 阶段 初步认识好 MMA 指令后,我们来看一下 flash-attn 要怎么用 MMA 指令来实现。 问题分析 首先,需要明确 flash-attn 的输入输出,那么在很多大模型的 prefill 阶段,输入输出 Q K V O 张量的输入张量维度都是:[batch_size, num_heads, seq_len, head_dim],我们为方便起见,
Ray 框架初步认知与理解
Ray 框架内部原理理解 本文参考 Ray 论文 Ray 框架解决的需求 最新的强化学习(RL)算法强调让 AI 通过与环境的不断交互来学习提升自己。其核心目标是要让 AI 学习一个策略,这个策略将根据环境的变化,让 AI 自主做出相关的应对动作,并在环境中不断提升应对策略,久而久之 AI 将学会一个有效的任务(例如赢得游戏或驾驶一架无人机)。在大规模的应用中,寻找有效的策略需要三大能力: 模
Tensor Core 和 MMA
Tensor Core 与 MMA 自 Volta 架构开始,nvidia 在显卡上装上了 Tensor Core 架构。该架构是为满足深度学习中所需的大量矩阵类运算需求而设计的硬件架构,专门提供高效小块的矩阵乘法:D=A×B+CD=A\times B+CD=A×B+C。 因为深度学习中高精度 float 运算不是必须的,所以 Tensor Core 还支持更低精度的计算,更低精度的计算意味着更高
ncompute 学习与使用
ncompute 使用示例 准备工作 机器设备:NVIDIA GeForce RTX 4060 laptop CUDA 环境:12.4 docker 环境:nvcr.io/nvidia/pytorch:24.05-py3 我简单地写了一个 elementwise 的 kernel 函数,并使用 ncu 对其做性能分析 1234567891011__global__ void elementwi
DeepSeek 背后的 MLA 和 MoE 架构
最近,DeepSeek(深度求索)公司推出的 DeepSeek-V3 和 DeepSeek-R1 大火,吸引了太平洋两岸所有关心关注 AI 发展的人的目光。本文试图从 DeepSeek 这轮爆火现象的背后,探究其中的架构创新,进而挖掘它如此低廉却好用的原因。 MLA(Multi-head Latent Attention) 一句话说明什么是 MLA 为了进一步解决 KV cache 在模型推理中的