# Phase 14: Flash Attention 2 for SM120 — Design Document

## Goal

用自写的 Flash Attention 2 CUDA kernel 替换 naive attention (Phase 5)。消除 O(S²) 显存分配，支持 GQA kernel 内部索引（消除 repeat_kv 开销）。

## 硬件约束: FA4 不适用于 RTX 5090

Flash Attention 已发展到第 4 代 (FA4, arxiv 2603.05451)，但各版本有明确硬件依赖：

| 版本 | 目标架构 | 关键硬件特性 | RTX 5090 (SM120) |
|------|---------|------------|-----------------|
| FA2 | 通用 CUDA (SM75+) | shared memory + HMMA | **兼容** |
| FA3 | Hopper SM90 (H100) | TMA + WGMMA + warp specialization | 不兼容 |
| FA4 | Blackwell SM100 (B200/B300) | TMEM + async MMA + 2-CTA mode | 不兼容 |

RTX 5090 使用消费级 Blackwell (GB202, SM120)，与数据中心 Blackwell (B200, SM100) 是不同硅片。SM120 **没有 TMEM (Tensor Memory)**，这是 FA4 kernel 设计的核心硬件依赖。这不是软件限制，是硬件级差异。

因此本项目实现 **FA2 算法**，使用标准 CUDA (shared memory + 标准 HMMA)。

## Naive Attention 的问题

Phase 5 的 naive attention 流程:
```
k_t = K.transpose(2,3).contiguous()    ← 分配 K^T 显存
scores = batched_matmul(Q, k_t)         ← 分配 [B,H,S,S] score 矩阵 (O(S²) 显存)
scores = scale(scores, 1/sqrt(d))       ← 逐元素 kernel
causal_mask(scores)                     ← 逐元素 kernel
weights = softmax(scores)               ← 分配 [B,H,S,S] weight 矩阵
output = batched_matmul(weights, V)     ← 最终结果
```

问题:
1. **显存 O(S²)**: score 和 weight 矩阵各需 `B × H × S × S × dtype_size`。S=2048, H=32, BF16 → 256 MB。S=8192 → 4 GB。
2. **GQA 预处理**: 在调用 attention 前需要 `repeat_kv_gpu` 将 K/V 从 8 heads 扩展到 32 heads，每层额外分配和拷贝。
3. **多次 kernel launch**: scale, mask, softmax 各一次 kernel launch + global memory round-trip。
4. **K^T materialization**: `K.transpose().contiguous()` 需要分配和拷贝。

## FA2 算法

核心思想: **不 materialize S×S 矩阵**。将 Q, K, V 分成 tiles，在 shared memory (SRAM) 中计算，使用 **online softmax trick** 边算边更新 running max 和 sum。

FA2 (Dao 2023) 相比 FA1 的改进: 外层循环遍历 Q tiles (而非 K/V)，减少 HBM 读写次数，提高并行性。

```
scale = 1 / sqrt(head_dim)

for each Q tile (q_start..q_start + BR):                   // 外层: Q tiles
    load Q_tile [BR, D] to shared memory (一次加载，内层复用)
    init per-row: O[D] = 0, m = -inf, l = 0

    for each K/V tile j (kv_start..kv_start + BC):          // 内层: K/V tiles
        // Causal tile-skip: 如果整个 K tile 在 Q tile "未来"，跳过
        if causal && kv_start > max_q_pos + kv_offset: skip

        load K_tile [BC, D] to shared memory
        S = Q_tile @ K_tile^T * scale                       // [BR, BC], in registers
        if causal: mask S[r][c] = -inf where kv_pos > q_pos

        // Online softmax update
        m_new = max(m, rowmax(S))
        P = exp(S - m_new)
        l_new = exp(m - m_new) * l + rowsum(P)
        O = exp(m - m_new) * O                              // rescale accumulator

        load V_tile [BC, D] to shared memory (复用 K 的空间)
        O += P @ V_tile                                     // accumulate

        m = m_new, l = l_new

    O = O / l                                               // final normalize
    write O[BR, D] to HBM (convert FP32 → BF16)
```

## 实现细节

### Kernel 配置

| 参数 | 值 | 说明 |
|------|---|------|
| BR (Q tile rows) | 64 | Q tile 大小 |
| BC (K/V tile rows) | 64 | K/V tile 大小 |
| head_dim | 运行时参数 (≤128) | 支持 64 (GPT-2) 和 128 (Qwen3) |
| Block size | 128 threads | 64 线程各 own 一行 Q，其余协助加载 |
| Grid | (q_tiles, batch × num_q_heads) | 每个 block 处理一个 Q tile + 一个 head |

### Shared Memory (BF16 存储)

```
smem_q [BR × head_dim] BF16 = 64 × 128 × 2 = 16 KB (加载一次，内层复用)
smem_kv[BC × head_dim] BF16 = 64 × 128 × 2 = 16 KB (K 和 V 交替使用)
────────────────────────────────────────────
Total: 32 KB (SM120 默认 48 KB，余量充足)
```

### 线程映射

- Thread 0..63: 各 own Q_tile 的一行。负责该行的全部计算：dot products、softmax、PV 累加。
- Thread 64..127: 协助 shared memory 加载 (K/V tile)，不参与计算。
- 加载模式: 每个 thread 加载 `(BR × head_dim) / 128 = 64` 个 BF16 元素。

### Per-Thread Register 使用

```
O_acc[128] FP32     = 512 bytes (128 regs)  — 输出累加器
P[64]      FP32     = 256 bytes (64 regs)   — 当前 tile 的 softmax 后权重
m, l       FP32     = 8 bytes (2 regs)      — online softmax running state
循环变量 + 临时      ≈ 16 regs
────────────────────────────────────────────
Total: ~210 regs/thread (max 255，在限制内)
```

### GQA 支持

每个 thread block 处理一个 Q head，通过 `kv_head = q_head / (num_q_heads / num_kv_heads)` 映射到对应的 KV head。K/V 的数据指针直接指向 KV head 的存储，无需 repeat_kv。

```
// 32 Q heads, 8 KV heads → heads_per_group = 4
// Q head 0,1,2,3 → KV head 0
// Q head 4,5,6,7 → KV head 1
// ...
kv_head = q_head / heads_per_group;
K_ptr = K + (batch * num_kv_heads + kv_head) * kv_len * head_dim;
```

### Causal Mask

两级优化:
1. **Tile-level skip**: 如果 `kv_tile_start > max_q_pos + kv_offset`，整个 K/V tile 都在未来，跳过（减少 ~50% 计算）。
2. **Element-level mask**: 在 tile 内部，`if kv_pos > q_pos + kv_offset: S = -inf`。

`kv_offset = kv_len - q_len` 处理 decode 时 KV cache 长于 Q 的情况。

## 与 Naive Attention 的对比

| 特性 | Naive (Phase 5) | FA2 (Phase 14) |
|------|----------------|----------------|
| 显存 | O(B × H × S²) | O(B × H × S × D) |
| GQA | 需要 repeat_kv (分配+拷贝) | Kernel 内部索引 (零开销) |
| K^T | 需要 transpose+contiguous | Kernel 内部计算 |
| Kernel launches | 6 (matmul, scale, mask, softmax, matmul, ...) | 1 (单个 fused kernel) |
| S=8192 可行性 | OOM (~4 GB score matrix) | 可行 (32 KB shared memory) |

## 源码结构

```
csrc/attention/flash_attention.cu     — FA2 kernel (BF16 in, FP32 accumulate, BF16 out)
crates/xserv-kernels/src/attention.rs — flash_attention() Rust wrapper + 原 attention() 保留
crates/xserv-model/src/qwen3.rs       — forward_gpu_cache 调用 flash_attention
```

## 已知局限与后续优化方向

1. **Decode (Q_len=1) 效率低**: BR=64 线程中只有 1 个 active（owns_row）。应写专用 decode attention kernel，沿 KV 维度 parallel reduction。
2. **无向量化加载**: 当前逐元素 bf16→f32 转换，应改用 `float4` 或 `__nv_bfloat162` 批量加载。
3. **Register tiling**: 每个 thread 目前串行计算 dot product (128 MADs per K column)。可改为多线程协作。
4. **K/V double buffering**: 可在计算当前 tile 时预加载下一个 tile 到另一半 shared memory。
5. **Tile size 调优**: 更大的 tile (BR=128) 可能在长 sequence 时更优，需要 opt-in shared memory。

## Test Plan

- [x] 正确性: logits 与 HF transformers 对比 (top-1 match 9/10, top-5 overlap 4.0/5)
- [x] 生成质量: 52/52 prompt 生成连贯文本，中英文均可
- [x] SSE streaming 正常工作
- [x] 性能: 12.9 tok/s (vs naive 10.3 tok/s, +25%)
- [ ] 长 sequence (S=4096, S=8192): 验证 naive OOM 而 FA2 正常
- [ ] ncu profile: compute utilization, memory throughput