矩阵乘法与Flash Attention优化笔记

一、分块矩阵乘法（Tiled Matrix Multiplication）

1. 核心原理

矩阵乘法基本式为（A为矩阵，B为矩阵，C为矩阵）。常规计算中输入元素需从全局内存重复读取多次，分块矩阵乘法将大矩阵A、B划分成大小为T的小“块（Tile）”，计算时：

• 外层循环：遍历这些块（紫色部分为外层循环的块，深蓝色是当前外层循环处理的块）。
• 内层循环：在每个块内部，遍历块里的元素（绿色部分为内层循环的元素，亮绿色是当前内层循环处理的元素）。
• 结果计算：通过分块计算逐步得到结果矩阵C的临时块（橙色部分），最终整合为完整的C。

2. 内存访问优化

二、burst section相关

1. 所属硬件

“burst section”通常指GPU（图形处理器）中的概念，代表内存突发传输中的一个数据段。GPU为提高内存访问效率，采用成块（burst）读取数据的方式，而非逐个字节读取。虽CPU在缓存读取等操作中也有成批读取数据的概念，但在矩阵分块优化及相关术语使用语境中，尤其涉及大规模并行计算和内存访问优化时，更多在GPU计算加速场景下讨论。

2. 内存对齐与数据加载

• 内存对齐（Aligned Layout）：当矩阵分块（tile）和内存突发传输段对齐时，GPU可高效加载数据，能一次性读取整个分块的数据，减少读取次数，提升性能（如“Aligned Layout”示例展示，快速加载“One Nice Tile”）。
• 内存未对齐（Unaligned Layout）：若矩阵分块和内存突发传输段不对齐，一个分块的数据可能分散在多个突发传输段中，GPU需多次读取不同突发传输段获取完整分块数据，导致性能下降（如“Unaligned Layout”示例，产生“Two Bad Tiles”，数据加载效率低）。

三、方阵矩阵乘法性能图表（Matrix mystery）

1. 坐标轴与核心指标

• 横轴：矩阵相关规模参数（如矩阵维度等）。
• 纵轴：TF/s（每秒万亿次浮点运算），衡量矩阵乘法计算性能。

2. 关键技术与趋势

• Compute Intensity（计算强度，粉色标注）：随矩阵规模等因素变化，计算强度提升推动性能（FLOPs）上升。计算强度为“计算操作与内存访问的比例”，比例越高，计算越“密集”，能更充分利用硬件计算能力。
• Tiling（分块，黄色标注）：经典优化手段，通过把大矩阵拆成小“块（Tile）”计算，减少内存访问开销、提升数据复用率，显著提高性能（黄色箭头指向区域性能因分块优化明显提升）。
• Wave Quantization（波量化，绿色标注）：优化思路（可能涉及数据量化、压缩等，减少计算或内存传输开销），绿色圈出区域性能在该优化下有特定变化趋势。

3. 整体意图

标题“Matrix mystery”（矩阵之谜）及文字“We understand some of this (compute intensity, tiling)… let’s take a closer look…”体现：虽已理解“计算强度”“分块”等部分优化作用，但仍需深入分析这些技术（及“波量化”这类技术）如何共同影响矩阵乘法性能，探索其中“奥秘”。

四、不同k值与内存对齐关系

1. 基础关系

矩阵乘法分块优化（Tiling）中，内存对齐是关键因素。矩阵分块与内存突发传输段（burst section）对齐时，数据加载效率高；不对齐时，数据加载效率低。

2. 不同k值影响

K通常代表矩阵分块的某个维度参数（如分块大小或划分粒度），不同K值改变矩阵分块尺寸和形状，影响其与内存突发传输段匹配关系：

• 合适的K值（利于对齐）：K取值使分块边界与内存突发传输段边界匹配（如分块尺寸是内存突发传输段大小的整数倍），加载分块数据时实现内存对齐，一次突发传输完整加载一个分块，数据加载高效。
• 不合适的K值（导致未对齐）：K取值不合适，分块边界与内存突发传输段边界不匹配，一个矩阵分块数据可能跨越多个内存突发传输段，加载一个分块需多次突发传输，内存未对齐，严重降低数据加载效率。

3. 图表体现

不同K值（K = 2、K = 8、K = 16、K = 32）对应FLOPs曲线走势不同。不合适的K值因导致内存未对齐，数据加载效率低，影响整体矩阵乘法计算性能，FLOPs表现不如内存对齐情况。

五、Softmax计算（普通与在线）

1. 普通Softmax（Normal softmax，左侧）

• 公式：，通过减去输入向量x中的最大值进行数值稳定（避免指数运算数值过大溢出），计算每个元素的Softmax结果。
• 算法（Algorithm 2: Safe softmax）：
  1. 初始化最大值为负无穷。
  2. 遍历输入向量x的每个元素，逐步计算当前最大值（找到整个向量的最大值）。
  3. 初始化总和为0。
  4. 遍历输入向量x的每个元素，计算并累加到（得到分母的总和）。
  5. 对每个元素，计算，得到Softmax结果。

2. 在线Softmax（Online softmax，右侧）

• 算法（Algorithm 3: Safe softmax with online normalizer calculation）：
  1. 初始化最大值为负无穷，总和为0。
  2. 遍历输入向量x的每个元素：
     • 逐步更新当前最大值。
     • 在线更新总和，利用前一次的最大值和当前最大值m_j调整累加项（），避免后续重复计算。
  3. 对每个元素，计算，得到Softmax结果。

3. 核心差异

普通Softmax“先找全局最大值，再统一计算分母总和”；在线Softmax“边遍历元素、边更新最大值，同时在线维护分母总和”，在处理大规模数据或对计算效率（尤其是内存访问模式）有要求的场景下，在线Softmax可能更具优势（如更友好的缓存利用、流式计算特性）。

六、Flash Attention优化（含Softmax与输出部分）

1. 整体优化方式

2. Softmax节省计算细节

Flash Attention节省计算非对应“O(2)”这类计算，主要通过：

• 避免全局softmax的归一化计算：传统Attention计算softmax需对整个注意力得分矩阵归一化（计算所有元素指数值，再求和得到归一化因子，计算复杂度高，序列长度n较大时计算量随增长）；Flash Attention通过分块计算和增量式归一化因子计算（如计算和并复用结果），避免全局softmax归一化计算，减少指数运算和求和运算次数。
• 分块矩阵乘法减少内存和计算开销：传统Attention计算注意力得分对完整矩阵相乘（规模大时内存占用大、计算成本高）；Flash Attention将键矩阵分块，查询矩阵分别与分块转置矩阵相乘，降低每次计算内存需求，适配GPU并行计算特性。
• 中间结果复用减少重复计算：传统Attention可能未充分复用中间结果导致重复计算；Flash Attention复用中间计算结果（如依赖），避免重复计算指数和累加值。
• 内存层级优化减少数据传输开销：传统Attention数据在不同内存层级间频繁传输，传输开销大；Flash Attention合理安排计算和存储位置（如在SRAM中计算部分结果），减少数据传输，间接节省计算时间。

3. 输出部分（和）节省计算

核心逻辑

输出O是注意力权重与值（Value）矩阵的加权和，Flash Attention中输出拆分为和两部分组合：是分块的归一化因子，、是分块的注意力权重，、是分块的值矩阵）。

节省计算关键

• 复用的结果：计算时直接复用之前计算的，不重新计算与相关的加权和，避免对和的重复运算，减少矩阵乘法和加权求和计算量。
• 增量式归一化（Rescaling）：通过的“重新缩放（Rescaling）”操作，将分块计算结果增量式合并，不对整个矩阵重新做全局softmax归一化，避免大规模矩阵重复归一化计算，节省指数运算和求和开销。

迭代方式优势

Flash Attention输出合并阶段采用增量式、复用中间结果的方式，处理后续tile（如、对应的分块）时，复用之前tile（、对应的分块）计算的中间结果（如、归一化因子等），通过“重新缩放（Rescaling）”将新tile计算的注意力权重与值的加权和和之前tile结果增量式合并，避免对每个tile“重新遍历、完整计算一遍softmax相关所有步骤”，节省大量重复计算。

4. 实际分块方式

分块依据

• 内存层级适配：GPU有不同内存层级（SRAM速度快但容量小，HBM容量大但速度相对慢），分块大小设计成能让分块数据“刚好适配高速缓存（如SRAM）”，计算单个分块时数据留在高速缓存，减少慢速全局内存访问。
• 并行计算效率：分块能让GPU的线程束（Warp）或线程块（Block）高效并行计算，分块维度匹配GPU线程并行度，使每个线程或线程束“负载均衡”处理分块内计算。

实际分块动态性

分块大小非固定“两块”“三块”，而是动态调整：根据输入序列长度、模型隐藏层维度等参数，计算最优分块尺寸（如每个分块的token数量、特征维度等）；例如长序列（几千甚至上万个token）会分成多个小分块，逐个处理，每个分块大小优化为“能让该分块的注意力计算在GPU上达到最高吞吐量”。

核心目的

让每个分块的中间结果（如注意力得分、softmax中间值）保存在高速内存中，避免频繁从低速全局内存读取/写入；让分块内的矩阵乘法、softmax等操作能被GPU高效并行执行，最大化FLOPS（每秒浮点运算次数）。

安装

brew install --cask tla+-toolbox

案例——判断死锁

------------------------------ MODULE deadlock ------------------------------

EXTENDS Integers
VARIABLES locks, locksA, locksB

Init == 
    /\ locks = {}
    /\ locksA = {}
    /\ locksB = {}

A_REQUIRE_LOCK_1 ==
    /\ ~(1 \in locks)
    /\ locks' = locks \cup {1}
    /\ locksA' = locksA \cup {1}
    /\ locksB' = locksB

A_REQUIRE_LOCK_2 ==
    /\ 1 \in locksA
    /\ ~(2 \in locks)
    /\ locks' = locks \cup {2}
    /\ locksA' = locksA \cup {2}
    /\ locksB' = locksB
    
B_REQUIRE_LOCK_1 ==
    /\ 2 \in locksB
    /\ ~(1 \in locks)
    /\ locks' = locks \cup {1}
    /\ locksB' = locksB \cup {1}
    /\ locksA' = locksA

B_REQUIRE_LOCK_2 ==
    /\ ~(2 \in locks)
    /\ locks' = locks \cup {2}
    /\ locksB' = locksB \cup {2}
    /\ locksA' = locksA

Next == A_REQUIRE_LOCK_1
    \/ A_REQUIRE_LOCK_2
    \/ A_REQUIRE_LOCK_1
    \/ A_REQUIRE_LOCK_2
    
Spec == Init /\ [][Next]_<<locks, locksA, locksB>>

=============================================================================
\* Modification History
\* Last modified Wed Aug 06 12:14:05 CST 2025 by lucas
\* Created Wed Aug 06 12:05:08 CST 2025 by lucas

本质上就是用 latex 的语法去写状态机的状态转移表达式，包括状态转移条件，以及下一个状态。注意每个状态都需要在状态转移表达式中声明（即便没有发生变化，如 locksA' = locksA）

成功检测到了死锁：

使用 std::tie 读取 std::pair

std::pair<int, double> p{1, 2.0};
int i;
double d;
std::tie(i, d) = p;

std::atomic::is_always_lock_free

判断是否当前类型可以使用无锁原子实现。

std::span

std::span 就是一个连续对象存储的观察者。类似std::string_view 是 std::string 的观察者。

int printSize(std::span<int> s) { std::cout << s.size() << std::endl; }

int main() {
  std::vector<int> v{1, 2, 3};
  int a[10];
  std::array<int, 30> arr;
  printSize(v); // 3
  printSize(a); // 10
  printSize(arr); // 30
}

std::vector 释放内存

std::vector::clear 并不会释放 capacity，需要使用 std::vector::shrink_to_fit 或：

std::vector<int> v{1, 2, 3};
std::vector<int>().swap(v);
std::cout << v.capacity() << std::endl; // 0

虚函数表占的空间

class A {
public:
    virtual void vfunc1();
    virtual void vfunc2();
    void func1();
    void func2();
private:
    int m_data1, m_data2;
};

手写一个 std::function

#include <iostream>

template <typename R, typename... Args> struct ICallable {
  virtual R invoke(Args &&...args) = 0;
  virtual ~ICallable() {}
};

template <typename T, typename R, typename... Args>
class CallableImpl : public ICallable<R, Args...> {
private:
  T callable;

public:
  explicit CallableImpl(T callable) : callable(callable) {}

  R invoke(Args &&...args) override {
    return callable(std::forward<Args>(args)...);
  }
};

template <typename Signature> class Function;

template <typename R, typename... Args> class Function<R(Args...)> {
private:
  std::unique_ptr<ICallable<R, Args...>> callable;

public:
  template <typename T>
  Function<R(Args...)>(T &&callable)
      : callable(std::make_unique<CallableImpl<T, R, Args...>>(
            std::forward<T>(callable))) {}

  Function(const Function<R(Args...)> &other) = delete;
  Function &operator=(const Function<R(Args...)> &other) = delete;

  R operator()(Args &&...args) {
    return callable->invoke(std::forward<Args>(args)...);
  }
};

int mmax(int a, int b) { return a > b ? a : b; }

int main() {
  Function<int(int, int)> f1 = mmax;
  std::cout << f1(3, 4) << std::endl;

  Function<int(int, int)> f2 = std::function<int(int, int)>(mmax);
  std::cout << f2(3, 4) << std::endl;

  Function<int(int, int)> f3([](int a, int b) { return a > b ? a : b; });
  std::cout << f3(3, 4) << std::endl;

  class Callable {
  public:
    int operator()(int x, int y) { return x > y ? x : y; }
  } c;
  Function<int(int, int)> f4(c);
  std::cout << f4(3, 4) << std::endl;
}

智能指针

std::make_shared 优于 std::shared_ptr<T>(new T) 的原因是：

1. 更简洁
2. std::make_shared 会分配一块连续的空间（控制块 + 数据块），locality 更好。而后者的内存空间是分散的。

手写 std::shared_ptr:

template <typename T> class ReferenceCount {
private:
  std::atomic<size_t> mRefCount;
  T *data;

public:
  ReferenceCount(T *data) : data(data), mRefCount(1) {}
  void increase() { mRefCount.fetch_add(1); }
  void decrease() {
    if (mRefCount.fetch_sub(1) == 1) {
      delete this;
    }
  }
  size_t count() const { return mRefCount.load(); }
};

template <typename T> class SharedPtr {
private:
  T *data;
  ReferenceCount<T> *mRefCount;

public:
  SharedPtr(T *data) : data(data), mRefCount(new ReferenceCount<T>(data)) {}

  SharedPtr(const SharedPtr<T> &other)
      : data(other.data), mRefCount(other.mRefCount) {
    mRefCount->increase();
  }

  SharedPtr<T> &operator=(const SharedPtr<T> &other) {
    if (mRefCount != other.mRefCount) {
      mRefCount->decrease();
      mRefCount = other.mRefCount;
      data = other.data;
    }
    mRefCount->increase();
    return *this;
  }

  size_t useCount() const { return mRefCount->count(); }

  T *operator->() { return data; }
};

template <typename T, typename... Args> SharedPtr<T> makeShared(Args... args) {
  return SharedPtr<T>(new T(std::forward<Args>(args)...));
}

int main() {
  SharedPtr<int> a = makeShared<int>(4);
  SharedPtr<int> e = makeShared<int>(5);
  std::cout << a.useCount() << std::endl;
  SharedPtr<int> b = a;
  std::cout << a.useCount() << std::endl;
  SharedPtr<int> c(a);
  std::cout << a.useCount() << std::endl;

  b = e;
  std::cout << a.useCount() << std::endl;
  std::cout << e.useCount() << std::endl;

  struct Person {
    Person(std::string name, int age) : name(name), age(age) {}
    std::string name;
    int age;
  };

  SharedPtr<Person> p = makeShared<Person>("Alice", 32);
  std::cout << p->name << ", " << p->age << std::endl;
}

输出：

1
2
3
2
2
Alice, 32

手写 std::any

老样子，使用类型擦除技术（基于特化）。这里的思路我总结下来是这样：
std::any 他是不需要指定类型的，比如：

std::any a = 123;

这里实际上只有在构造的时候才会用到类型，对应一个特化的构造函数：

template<typename T>
Any(T val);

例如：

std::any a = 123;
std::any a = std::string("Hello");

将对应两个特化的构造函数。

既然 std::any 不需要指明其存储的类型 T，那么他的所有成员变量一定是和类型 T 无关的，但是我们总要有个地方存这个 T 类型的 value。

所以可以想到，我们定义一个接口，该接口和 T 类型无关，其实现类可以与 T 有关，该实现类是真正存储 T 类型 value 的地方。而 std::any 则可以存储这个接口。

struct IValueHolder {
  virtual const std::type_info &type() const = 0;
};

template <typename T> class ValueHolder : public IValueHolder {
private:
  T val;

public:
  explicit ValueHolder(T val) : val(val) {}

  const std::type_info &type() const override { return typeid(val); }

  inline const T &value() const { return val; }
};

class Any {
private:
  std::unique_ptr<IValueHolder> holder;

public:
  template <typename T>
  Any(T val) : holder(std::make_unique<ValueHolder<T>>(val)) {}

  template <typename T> T AnyCast() {
    if (typeid(T) != holder->type()) {
      throw std::bad_cast();
    }
    return static_cast<ValueHolder<T> *>(holder.get())->value();
  }
};

int main() {
  Any a = 123;
  std::cout << a.AnyCast<int>() << endl;
  std::cout << a.AnyCast<double>() << endl;
}

输出：

123
libc++abi: terminating due to uncaught exception of type std::bad_cast: std::bad_cast
[1]    25963 abort      /Users/lucas/projects/c++/bin/my_test

enum class

1. 避免命名空间冲突。
2. 可以指定底层存储的类型，int or long long…
3. 可以前置声明。

std::atomic

std::atomic 背后的缓存行锁定（Cache Line Locking）是实现原子操作的关键机制，它依赖于现代 CPU 的缓存一致性协议（如 MESI）和总线锁定（Bus Locking）。当线程对 std::atomic 变量执行原子操作时，CPU 首先尝试通过缓存行锁定实现高效并发：如果该变量所在的缓存行被当前 CPU 独占（Modified 状态），则直接修改并标记为脏；若被其他 CPU 共享（Shared 状态），则通过总线发送 invalidate 信号，迫使其他 CPU 放弃缓存行的所有权，从而获取独占权后执行操作。若缓存行锁定失败（如跨缓存行的操作），CPU 会退化为总线锁定，临时阻塞整个内存总线，确保操作的原子性。这种分级机制平衡了性能与正确性，使 std::atomic 在大多数场景下仅需短暂的缓存行竞争，避免了昂贵的总线锁定开销。

std::atomic<int> flag{0};
int data = 0;

// 线程 A（写端）
data = 42;                          // 非原子写入
flag.store(1, std::memory_order_release); // release 保证 data 的写入不会重排到之后

// 线程 B（读端）
if (flag.load(std::memory_order_acquire) == 1) { // acquire 保证读到 release 的写入
    assert(data == 42); // 一定成功！data 的写入对线程 B 可见
}

std::memory_order_release 保证它之前的所有读写都不会重排到它后面，而 std::memory_order_acquire 又保证了它之后的所有读写不会重排到它前面，这样就形成了 release-acquire 模式，前者保证 flag.Store 之前 data = 42 已经写入；后者保证 flag.load 之后，data = 42 才被读取。

noexcept

std::vector 扩缩容的时候，如果元素的移动构造函数被标记为 noexcept，那它就会大胆地调用移动构造函数，否则就会使用相对更安全的拷贝构造函数。std::sort 的优化类似。

class A {
public:
  A() = default;
  A(const A &) { std::cout << "copy" << std::endl; }

  A(A &&) { std::cout << "move" << std::endl; }
};

int main() {
  std::vector<A> aVec;
  aVec.emplace_back();
  aVec.emplace_back();
}

输出：

copy

将移动构造函数标记为 noexcept，输出：

move

std::enable_shared_from_this

有些时候，我们需要在一个 std::shared_ptr 托管的类的成员函数中，创建它的一个 shared_ptr 新实例，传给某些函数，例如：

class Manager;
class Component {
private:
  int *value = new int(3);
  Manager *mgr;

public:
  Component(Manager *mgr) : mgr(mgr) {}
  ~Component() {
    std::cout << "Deconstruct" << std::endl;
    delete value;
  }
  void DoSomething();
  friend class Manager;
};

class Manager {
public:
  void Modify(std::shared_ptr<Component> c) { *c->value = 3; }
};

void Component::DoSomething() { mgr->Modify(std::shared_ptr<Component>(this)); }

int main() {
  Manager mgr;
  auto c = std::make_shared<Component>(&mgr);
  c->DoSomething();
}

上述的代码会导致 double free 的问题，正确做法如下：

class Component : public std::enable_shared_from_this<Component> {
  ...
}
void Component::DoSomething() { mgr->Modify(shared_from_this()); }

其原理是：构造 Component 的 shared_ptr 的时候，会同时定义一个 weak_ptr 指向 shared_ptr。在调用 shared_from_this 的时候，会由该 weak_ptr 创建对应的 shared_ptr。

注意，此处的 std::enable_shared_from_this<Component> 必须是 public 继承，不然创建 std::shared_ptr 的时候无法正确访问 enable_shared_from_this 中的 weak_ptr。

std::deque 实现原理

std::deque 通常采用 分块存储（Chunked Storage） 的方式实现，即：
由多个固定大小的 “块（Chunks）” 组成，每个块是一个连续的内存数组（称为 “缓冲区” 或 “页”）。

使用一个中央控制结构（如指针数组或动态数组）管理这些块，记录每个块的地址。
这种设计使得：
头尾插入/删除高效（不需要像 vector 那样整体移动元素）。
随机访问（operator[]）比 list 快（计算块索引 + 块内偏移）。

典型的 std::deque 实现（如 GCC 的 libstdc++）采用 “Map + Chunks” 结构：

+-------------------+     +-------------------+     +-------------------+
| Chunk 0 (front)   | --> | Chunk 1           | --> | Chunk N (back)    |
+-------------------+     +-------------------+     +-------------------+
| [0] | [1] | ...   |     | [0] | [1] | ...   |     | [0] | [1] | ...   |
+-------------------+     +-------------------+     +-------------------+

• Map（控制中心）：一个动态数组（如 T**），存储指向各个块的指针。
• Chunk（块）：每个块是一个固定大小的数组（如 512B 或 4KB），存储实际数据。

1. PVC（Persistent Volume Claim，持久卷声明）：是用户对存储资源的请求，用于向 Kubernetes 申请存储，抽象化存储的使用细节，方便应用获取和使用持久化存储。
2. PV（Persistent Volume，持久卷）：是 Kubernetes 中底层存储的抽象，代表具体的存储资源（如 NFS、云硬盘等），通过定义存储的类型、容量、访问模式等参数，供 PVC 动态绑定调用。
3. StorageClass：用于定义存储资源的配置模板，可动态创建 PV，支持按需提供存储服务，通过参数配置存储的类型（如 SSD、HDD）、备份策略等，实现存储资源的自动化管理和动态供给。

静态绑定

以使用NAS静态存储卷为例：

• PV

  静态绑定的方式需要手动配置 PV，并设定 labels：

  apiVersion: v1
  kind: PersistentVolume
  metadata:
  name: pv-nas
  labels:
      alicloud-pvname: pv-nas
  spec:
  capacity:
      storage: 5Gi
  accessModes:
      - ReadWriteMany
  csi:
      driver: nasplugin.csi.alibabacloud.com
      volumeHandle: pv-nas   # 必须与PV Name保持一致。
      volumeAttributes:
      server: "0c47****-mpk25.cn-shenzhen.nas.aliyuncs.com"  # NAS挂载点，与集群VPC一致。
      path: "/csi"  # 挂载子目录。
  mountOptions:
  - nolock,tcp,noresvport
  - vers=3
  

• PVC

  PVC 通过 label selector 绑定上面手动配置的 PV。

动态绑定

1. 用户创建 PVC：用户通过 PVC 声明存储需求（如容量、访问模式、存储类型等），但此时无对应的 PV 绑定。
2. StorageClass 匹配与创建：PVC 会根据指定的 StorageClass （或默认类），触发存储插件（如云厂商存储插件、NFS 插件等）动态创建符合要求的 PV，并将两者绑定。
3. 应用使用存储：绑定完成后，Pod 可通过 PVC 挂载自动创建的 PV，实现按需获取存储资源，无需手动预先创建 PV。

PVC 和 PV 以 storageClassName 为桥梁，以 StorageClass 为中介动态绑定。若 storageClassName 为空，则 PVC 将会和默认的 PV 绑定。

StorageClass 的 allowVolumeExpansion 属性用于表明该存储是否允许扩容卷，以阿里云 NAS为例，若该属性为 true，则可以通过 patch PVC，修改其声明的卷大小来动态扩缩容。

CSI Driver

Understanding CSI architecture and communication

• Daemonset Pod 应用可在所有节点上运行。
• StatefulSet / Deployment Pod 应用作为 Kubernetes controller 运行。
• 在 Daemonset Pod中，CSI Driver 通过 gRPC 和 socket 与 Kubelet 通信。
• node-driver-registrar 将 CSI Driver 作为插件注册到 Kubelet。
  https://speakerdeck.com/bells17/kubernetes-and-csi?slide=44
• 在 StatefulSet / Deployment Pod 中，CSI Driver 通过 gRPC 和 socket 与 CSI sidecar 应用通信。
• CSI sidecar 应用作为 Kubernetes controller 运行。
• 例如，创建 PVC 资源时，external-provisioner 会监听该事件，并调用 CSI Driver 的 CreateVolume rpc。
  https://speakerdeck.com/bells17/kubernetes-and-csi?slide=47

标量量化（Scalar Quantization）

FP32 -> FP16，FP16 -> INT8

以 FP16 -> INT8 为例，将 FP16 的数值映射为 INT8 的范围。如果 tensor 中存在离群值（远大于平均的值），在 de-Quantized 之后可能导致其他值消失（接近 0）。解决方法是把这部分离群值拎出来，不量化，按照原有精度处理。

乘积量化（Production Quantization）

将原向量拆分成若干子向量，找到子向量距离最近的 cluster（聚类）对应的标识。由这些标识构成量化后的向量。

以上图为例，假设绿色为 0，蓝色为 1，红色为 2，则量化后的向量为：
[0, 1, 1, 2]。

二进制量化与 RaBitQ

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

二进制量化直接将向量的每一维度映射到一个 bit（0 或 1），常见的做法就是将大于 0 的维度设置为 1，否则则为 0。然而这种做法虽然极大地压缩了空间，但是也会显著影响召回率。

RabitQ 通过将向量进行归一化，将其放入一个单位高维球中，这个单位高维球中有一个 D 维的超正方体（D 为向量维度）。靠近哪个顶点，就被划定为这个顶点，以这个顶点对应的向量作为量化结果。

References

• RocksDB 源码 inlineskiplist.h

SkipList::Node 的实现

RocksDB 分配了一块连续的空间来存储节点在不同 level 的 next 指针。同时，使用 std::atomic 保证可以在无锁的情况下更新链表。其伪代码如下所示，可以说是把内存空间利用得严严实实。

#include <atomic>
#include <cstring>
#include <memory>

// https://github.com/facebook/rocksdb/blob/ef67339175c1aebe039c0671560b608e22656612/memtable/inlineskiplist.h
class SkipList {
public:
  struct Node {
  private:
    std::atomic<Node *> next_[1];

  public:
    void StashHeight(size_t h) { memcpy(&next_[0], &h, sizeof(h)); }

    size_t UnstashHeight() {
      size_t rv;
      memcpy(&rv, &next_[0], sizeof(rv));
      return rv;
    }

    Node *Next(size_t level) {
      return (&next_[0] - level)->load(std::memory_order_acquire);
    }
  };

  Node *AllocateNode(size_t height, size_t key_size) {
    size_t prefix = sizeof(std::atomic<Node *>) * height - 1;
    char *raw = new char[prefix + sizeof(Node) + key_size];
    Node *node = reinterpret_cast<Node *>(raw + prefix);

    node->StashHeight(height);

    return node;
  }
};

此处有个很有意思的点，在 AllocateNode 生成一个随机高度的节点的时候，会把他的高度通过 StashHeight 的方式暂时存储在 next_[0] 中。在后续的操作中，可以通过 UnstashHeight 读取存储的高度，并进行后续的 next 指针的更新（毕竟这些更新需要知道高度），可以说是不浪费任何内存空间。

其内存布局大概如下图所示（此处我做了简化，&next_[0] - n 其实就是相当于 next_[-n]，但是数组不支持负数）：

References

• [FIXME][EP09] 向量数据库和出海创业
• Locality Sensitive Hashing (LSH): The Illustrated Guide
• Write You a Vector Database

常用的向量检索算法

1. 量化（以一定的模型性能换取查询效率），RabitQ。

2. Locality Sensitive Hashing (LSH)：传统的 hash 函数一般都希望减少碰撞，将一个 key 映射到尽可能分散的空间，而 LSH 则恰恰相反，通过将比较近的 embedding 映射到尽可能相同的 bucket 中，这样就能快速找到相邻的节点。

   

3. NSW(Navigate Small World)

   

   选择若干个 entry，然后维护一个有固定大小上限的队列，以类似 BFS 的方式，将距离目标向量最近的向量入队。需要多个 entry 的原因是，如果只采用一个，很容易陷入局部最优解。

   为了加速查询，我们会限制图中点之间的边的数量为 max_m。当我们插入一个节点的时候，其邻居节点的边可能会超过 max_m，所以我们需要重新计算这些节点的最近邻。

4. HNSW(Hierarchical Navigate Small World)

   将图以类似跳表的形式组织，但是无论是 NSW 还是 HNSW 都会涉及大量磁盘随机 I/O，性能受限。

5. IVF(Inverted File)

   

   使用 Kmeans 将向量分为多个 cluster。并根据 cluster 的标识建立对向量的倒排索引。由于索引存储在一起，可以被顺序读写。

   在查询的时候，除了判断距离查询向量最近的质心对应的集群，还需要判断临近的集群。因为集群的质心可能距离查询向量比较远，但是边缘的点可能比较近。

   

References

1. 大模型推理加速与KV Cache（一）：什么是KV Cache
2. [FIXME][EP02][直播版] vllm源码讲解，分布式推理
3. [FIXME][EP05][直播版] vllm从开源到部署，Prefix Caching和开源答疑
4. CacheBlend-高效提高KVCache复用性的方法

Distributed Inference

具体的通信代码可以参考：vllm/model_executor/models/llama.py

分布式通信手段

1. nvlink
2. InfiniBand
3. RDMA

并行方式

1. TP(Tensor Parallel)：

   将 tensor 拆分运算，然后通过 all-reduce / all-gather 汇总。例如 LLM 的 Decoder 采用多头注意力（multihead-attention），假设有 个头部， 张卡，则可以每张卡 个头部并行计算。

   

   # llama前向传播的代码
   def forward(
       self,
       positions: torch.Tensor,
       hidden_states: torch.Tensor,
   ) -> torch.Tensor:
       qkv, _ = self.qkv_proj(hidden_states)
       q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
       q, k = self.rotary_emb(positions, q, k)
       attn_output = self.attn(q, k, v)
       output, _ = self.o_proj(attn_output)
       return output
   

2. PP(Pipeline Parallel)：

   一个 request 由多个 node 流水线式地执行。可以提高吞吐，但是无法降低 TTFT。

   每个worker负责一个layers的子集

   • vllm/model_executor/models/llama.py 中 self.start_layer –> self.end_layer
   • 在 worker 之间: communicate IntermediateTensor
   • vllm/worker/model_runner.py: 搜索 get_pp_group()

3. EP(Expert Parallel)：

   用于 MoE（Mixture of Expert）架构。在处理一个 request 的时候，只有一个很小的专家子集会被激活（通过门控控制）。router 会将 request 转发到相应的 expert 节点（MoE 的前置计算可能不是在相应的 expert 节点上进行的）。这里多个 expert 是并行的，同时计算不同的 batch。

   • Shuffle (DeepEP communication kernel)
   • Forward
   • Shuffle back

4. DP(Data Parallel)

   因为大型分布式生产环境中，EP >> TP（attention head 可能就十几个，而 experts 可能有几百个），这时候就需要 data parallel。

   • TP * DP == EP（通过请求并行的方式去拉满计算资源）
   • 在实践中难以应用。
   • 对请求进行padding避免造成死锁。

Prefix Caching

预备知识：

1. LLM 的 Prefill 阶段：当用户输入一个 prompt（例如 “Hello, how are”），模型会一次性处理整个输入文本，计算所有输入 token 的隐藏状态。这个阶段通常是并行计算的，效率较高；LLM 的 Decode 阶段：模型从第一个生成的 token 开始，逐个预测后续的 token（例如 “you”、”today” 等）。每个新 token 的生成都依赖于前一个 token。
2. 在多轮对话中，需要传递上下文，这些上下文总是共享一个前缀，如上图所示。显然，我们可以缓存这些前缀对应的 KV tensor，这也就是 Prefix Caching，缓存的对象便是 K 和 V tensor。
3. Prefix Caching 只节省了 Prefill 阶段的耗时（也就是降低了 TTFT ，Time To First Token），并不能节省解码阶段的耗时（也就是 TPOT ，Time Per Output Token）。

vllm 中的 Prefix Caching 主要思路：

1. 将若干个 tokens（默认 16 个）组成一个 block，对于每个 block，计算其 hash 值并更新 prefix hash。
2. 将 prefix hash -> kv cache 的映射存储到 kv store 中（例如 redis，注意后者的 kv 指的是键值存储）。
3. 若无法 allocate 缓存空间，以一定的策略（如 LFU，LRU）evict 一些缓存。

伪代码：

prefix_hash = ""
for chunk in chunked_token:
    chunk_hash = hash(prefix_hash + chunk)
    prefix_hash = chunk_hash

for chunk_hash, chunk_kv in zip(...):
    kv_store.put(chunk_hash, chunk_kv)

Prefix Caching 的扩展——CacheBalend

• Prefix Caching 只利用了前缀，缓存利用率有限；
• Full KV Cache 使用所有的 Cache，会忽视 cross attention，产生低质量结果；
• CacheBlend 通过重新计算一部分 KV，进行折中。

以下是忽视 cross attention 的结果：

PD 分离

为什么需要 PD 分离？

在大模型推理中，常用以下两项指标评估性能：
TTFT（Time-To-First-Token）：首 token 的生成时间，主要衡量 Prefill 阶段性能。
TPOT（Time-Per-Output-Token）：生成每个 token 的时间，主要衡量 Decode 阶段性能。
当 Prefill 和 Decode 在同一块 GPU 上运行时，由于两阶段的计算特性差异（Prefill 是计算密集型，而 Decode 是存储密集型），资源争抢会导致 TTFT 和 TPOT 之间的权衡。例如：

若优先处理 Prefill 阶段以降低 TTFT，Decode 阶段的性能（TPOT）可能下降。
若尽量提升 TPOT，则会增加 Prefill 请求的等待时间，导致 TTFT 上升。
PD 分离式架构的提出正是为了打破这一矛盾。通过将 Prefill 和 Decode 分离运行，可以针对不同阶段的特性独立优化资源分配，从而在降低首 token 延迟的同时提高整体吞吐量。

Prefill 和 Decode 阶段分别受限于什么？

Prefill 阶段：吞吐量随 batch size 增加逐渐趋于平稳。这是因为 Prefill 的计算受限特性（compute-bound），当 batch 中的总 token 数超过某个阈值时，计算资源成为瓶颈。

Decode 阶段：吞吐量随 batch size 增加显著提升。由于 Decode 阶段的存储受限特性（memory-bound），增大 batch size 可提高计算效率，从而显著增加吞吐量。

摘自 LLM推理优化 - Prefill-Decode分离式推理架构

1. Chunked Prefill
2. xP yD 问题，Prefill Instance 和 Decode Instance 的数量如何调整？
3. 先 D 后 P 的模式：D node收到后先检查是否有 KVCache，没有的话再转给 P node去做，这个思路主要考虑的是 TTFT。

传递 KV Cache

• 两种模式：pooling 模式，P2P 模式，LMCache都支持上面两种模式，Mooncake(pooling)，NIXL(p2p)。

• 怎么从 vllm 提取（注入）KVCache

  • connector API in vllm/worker/model_runner.py。
  • 在模型 forward 前：尝试接收 KVCache 并注入到到 vllm 的 pages memory 中。
  • 在模型 forward 后，将 KVCache 从 pages memory 中并将它发送出去。

  

• vllm/distributed/kv_transfer/kv_transfer_agent.py

  # 本质上是根据 model input 计算出 KVCache 放在 page memory中的什么地方
  def recv_kv_caches_and_hidden_states(
      self, model_executable: torch.nn.Module,
      model_input: "ModelInputForGPUWithSamplingMetadata",
      kv_caches: List[torch.Tensor]
  ) -> Tuple[Union[torch.Tensor, IntermediateTensors], bool,
          "ModelInputForGPUWithSamplingMetadata"]:
              
      return self.connector.recv_kv_caches_and_hidden_states(
          model_executable, model_input, kv_caches)
  

  可以先看 vllm/distributed/kv_transfer/kv_connector/simple_connector.py。

Speculative Decoding

前文提到 Prefill 是 gpu-bound，计算密集型（把 request 的 tokens 全部输入 llm，生成第一个 token，并构建 KVCache）；而 Decode 是 memory-bound，依赖 Prefill 阶段生成的 KVCache，访存的时间往往大于计算的时间。那么有没有一种方法可以在不怎么增加 memory access 的前提下，提升计算的吞吐呢？有的，兄弟，有的，Speculative Decoding。

Speculative Decoding 干了什么？其实就是去根据输入猜接下来的若干个 tokens 是什么，然后并行地进行验证。假如猜测 3 个 tokens，我们并行地验证这 3 个 tokens 是否正确的，由于是并行的，我们差不多只花费了原来只生成 1 个 token 的时间，最终获得了 3 个 tokens，也就是将吞吐提升了 3 倍，而访存只增加了（3 - 1）* token size。

那么如何猜呢？其实用一个古老的方法，n-gram 就行了。

An n-gram is a sequence of n adjacent symbols in particular order. The symbols may be n adjacent letters (including punctuation marks and blanks), syllables, or rarely whole words found in a language dataset; or adjacent phonemes extracted from a speech-recording dataset, or adjacent base pairs extracted from a genome. They are collected from a text corpus or speech corpus.

也就是说，我们根据一定的前缀，就能大致猜出后续的 token 搭配。Speculative Decoding 会根据输入构建 ngram，然后猜测后续的 tokens。

为什么这是有效的呢？下面几个 workload 就能说明这个问题：

1. 全文搜索，在用户给定的内容中寻找内容，或者给出一定答案。此处的回答一定是和用户内容强相关的，本质上会复读一部分；
2. 代码生成场景，变量名、函数名等都很容易被 ngram 预测。

Tree Verification

Model-based（draft model）Speculative Decoding

1. Parallel guessing（并行猜测）

   • 优点：快，在不知道第一个 token 情况下直接猜第二个。
   • 缺点：在猜测第二个 token 的时候不知道第一个token是什么，容易胡言乱语

2. Autoregression guessing（自回归猜测）

   • 优点：在猜测第二个 token 的时候知道第一个 token，准确率较高
   • 缺点：慢

RocksDB 架构

先写入 WAL，再写入 MemTable。MemTable 达到一定的大小就会变为 Immutable MemTable，后台异步以 SSTable 的形式 flush 到磁盘中。

MemTable 中还包含了一个 Bloom Filter，用于快速判断一个元素是否不存在。

Column Family

RocksDB 中列族是一个逻辑上的分组，列族之间使用不同的 LSM Tree，但是共享 WAL。

Version

RocksDB Version管理概述

Version 会指向不同的 SST，并且 Version 和 SSTable 都是以引用计数的方式维护，一旦引用计数变为 0，就会进行回收。

Write Batch

参考文献：

1. RocksDB 源码分析 – Write Batch
2. RocksDB 剖析 | WriteThread 如何控制并发写入流程

上面这篇文章讲的挺清楚的，RocksDB 默认是不开启并行写的，因为跳表结构并不适合并发写入，WAL 如果要并发写也需要一定同步机制维护其写入顺序。所以，RocksDB 采用的是分批写入的方式，以保证其吞吐量。Write Batch Group 会有一个 leader（当前的第一个 writer），然后后面的 writer 将成为 follower，构成一条双向链表（通过 link_older，link_newer）相连。这个 group 会有个大小上限，超过之后就会截断。截断后的新 writer 将会变为新的 leader，构成新的 group。

RocksDB 使用 atomic::<Writer *> 的形式组织链表，可以通过 CAS 以 latch-free 的方式添加 writer：

Writer* writers = newest_writer->load(std::memory_order_relaxed);
while (true) {
  w->link_older = writers;
  if (newest_writer->compare_exchange_weak(writers, w)) {
    return (writers == nullptr);
  }
}

这边正好学习了一下 std::atomic 的用法，比起其他语言，c++ 的原子对象的 load 操作还可以选择不同的参数：

1. std::memory_order_relaxed：仅保证原子性，不强制内存顺序。
2. std::memory_order_acquire：确保后续读操作不会重排到此操作之前。
3. std::memory_order_seq_cst（默认）：顺序一致性，提供最强的内存屏障。

尝试写了一个链表插入练习：

#include <atomic>
#include <iostream>
#include <thread>

struct ListNode {
  int val;

  std::atomic<ListNode *> next;
};

void insertListNode(std::atomic<ListNode *> &tail, ListNode **root, int val) {
  ListNode *newNode = new ListNode;
  newNode->val = val;
  newNode->next = nullptr;
  ListNode *oldTail = tail.load(std::memory_order_relaxed);
  while (true) {
    if (tail.compare_exchange_weak(oldTail, newNode)) {
      break;
    }
  }
  if (oldTail == nullptr) {
    *root = newNode;
  } else {
    oldTail->next = newNode;
  }
}

void printList(ListNode *root) {
  if (root == nullptr) {
    std::cout << "List is empty" << std::endl;
    return;
  }
  ListNode *p = root;
  while (p) {
    std::cout << p->val << " ";
    p = p->next.load(std::memory_order_relaxed);
  }
  std::cout << std::endl;
}

int main() {
  std::atomic<ListNode *> tail = nullptr;
  ListNode *root = nullptr;
  for (int i = 0; i < 3; i++) {
    std::thread t([&, i]() {
      while (1) {
        insertListNode(tail, &root, i);
        std::this_thread::sleep_for(std::chrono::milliseconds(300));
      }
    });
    t.detach();
  }

  while (1) {
    printList(root);
    std::this_thread::sleep_for(std::chrono::milliseconds(300));
  }
}

一种 Lazy Initialize 的方式：

std::optional<std::mutex> mtx;
mtx.emplace();

在线讲解：[Paper Reading] FoundationDB：开源分布式 KV 存储的内部实现原理

最近看 DeepSeek 的 3FS 的时候发现他们会用 FoundationDB 做 chunk metadata 的存储（3FS 有点像 GFS，也是以 chunk 作为存储单位）。正好不久之后要去阿里云开始一段存储相关的实习，学习一下相关知识。

FoundationDB 是一个分布式 KV 存储数据库，适用于读多写少的场景，支持 SSI（Serializable Snapshot Isolation）级别的事务。

FoundationDB 如何实现 SSI

SSI 这一级别是通过 OCC + MVCC 实现的。

首先，让我们复习一下事务隔离级别，可以从这篇文章切入：事务隔离级别备忘。通过 MVCC 的快照，我们可以很容易地实现 Snapshot Isolation，虽然这解决了脏读和幻读的问题，但是无法解决两种问题（也就是上面链接中的 P4 和 A5B 问题）。

• P4 Lost Update

  

  T2 修改 x 为 120，然而 T2 提交后被 T1 提交的修改覆盖，最终 x = 130。对外界而言，x = 120 这一修改不可见，被覆盖。

• A5B Write Skew

  

  x + y 应该满足约束 x + y <= 100，由于 T1 和 T2 的 write set 没有相交，所以两个事务都能提交成功，然而约束已经被破坏。

为了解决上述问题，FoundationDB 会维护一个 read set，如果 read set 中的某个 key range 对应的最新的一次 commit version 大于事务的 read version，就会 abort 当前事务。这种乐观的方式有若干好处：

• 无需锁，高效；
• 读多写少，abort 是小概率事件。

在 FoundationDB 中，resolvers 会保留最近一段时间（如 5 秒，MVCC window）内已提交的写入操作在内存中，这个时间范围就可以看作是一个 “窗口”。当一个新的事务进行冲突检测时，会将其读取的数据与这个窗口内的写入操作进行比较，以确定是否存在数据冲突。如果在这个窗口内，事务读取的数据被其他事务修改了，那么就会检测到冲突，该事务可能会被中止或重试。

FoundationDB 的架构

整个系统是自上而下启动的。FoundationDB 的控制面是通过 Paxos 算法选举出 Cluster Controller（即 leader），如下图所示。

下图为数据面，主要由 TS + LS + SS 组成。

TS

• Sequencer：
  • 负责获取 read/commit version。
• Proxy：
  • 一个中介，负责获取 read version 和事务提交。
• Resolver：
  • 分布式地解决读写冲突。
  • 多节点，按照 key range 分片（shard）。

LS

• 按照 key range 分片（shard）。

SS

• 单机存储系统，从 LS 的 WAL 异步复制数据。
• 多版本数据只存储在内存中
• 5s 的 MVCC window（太长了会产生大量假阳性）

Log System

• Log Server 下面绑了若干 Storage Server。
• 按照 key range 映射到若干个 replica。
• 之所以需要向未被映射到的 Log Server 传递空消息，是因为 Log Server 是按照顺序执行的，不这样做会产生“空洞”。

Commit Path

1. Client:
   1. Buffer Modifications
   2. Send Read/Write Set To Proxy
2. Proxy: Get Commit Version from Sequencer
3. Proxy: Send User Data KeyRange To Resolvers
4. Resolver: Check RW Conflict
5. Proxy: Commit Phase When no Conflict
   1. Send Data to LS
   2. Committed When All LS return Success
   3. Send Commit Version to Sequencer for Read
   4. Proxy return OK to Client
6. SS pull WAL from LS

Read: 1 roundtrip
Write: 4 roundtrip

Recovery && Replication

k = f + 1

这里 k 表示需要复制的日志份数，f 表示系统可容忍的故障节点数。即需要复制的日志份数比可容忍的故障节点数多 1，通过这种方式来保证即使部分节点故障，日志信息也不会丢失，能维持系统的正常运行和数据一致性。