deepseek-ai / deepseek-ai/FlashMLA

来自deepseek r1分析的进一步优化思路,大家看对不对 O(∩_∩)O

Open
#26 9 comments 91 reactions 0 assignees View on GitHub
Dominant language
C++
Stars
12.9k
Forks
1.2k
Avg merge
4h 20m
Merged PRs (30d)
2

Description

1. **Hopper异步拷贝强化**
```csrc/flash_fwd_mla_kernel.h
// 修改前
cute::cp_async_fence();

// 修改后
constexpr int kCpAsyncCount = 4; // 利用Hopper单周期4条cp.async
#pragma unroll
for (int i = 0; i < kCpAsyncCount; ++i) {
cute::cp_async<0x80, kCpAsyncCount>(tKgK_ptr + i, tKsK_ptr + i);
}
cute::cp_async_fence();
```
利用Hopper单周期发射4条异步拷贝指令的特性,提升SMEM填充吞吐量

2. **共享内存Bank冲突优化**
```csrc/flash_fwd_mla_kernel.h
// 修改前
using SmemLayoutP = Layout, Int, _1, Int>>;

// 修改后
using SmemLayoutP = Layout, Int<64>>, Stride<_64, _1>>; // 64字节对齐访问
```
调整P矩阵布局,确保每线程访问64字节对齐的连续内存,减少bank冲突

3. **动态指令流优化**
```csrc/flash_fwd_mla_kernel.h
// 修改前
if constexpr (!Is_causal) {
// 非因果分支
} else {
// 因果分支
}

// 修改后
constexpr int kCausalMask = Is_causal ? 1 : 0;
__builtin_assume(kCausalMask == 0 || kCausalMask == 1); // 帮助编译器优化分支
```
通过__builtin_assume提示编译器优化条件分支

4. **寄存器压力优化**
```csrc/flash_fwd_mla_kernel.h
// 修改前
Tensor tSrQ = ...;
Tensor tSrK = ...;
Tensor tSrS = ...;

// 修改后
union {
Tensor tSrQ;
Tensor tSrK;
Tensor tSrS;
}; // 共享寄存器空间
```
对临时张量使用联合存储,减少寄存器占用

5. **TMA加速全局内存访问**
```csrc/flash_fwd_mla_kernel.h
// 新增TMA描述符
cute::TmaDescriptor tma_q_desc, tma_k_desc;
cute::make_tma_copy_desc(
gQ, sQ, tma_q_desc,
Int{}); // 128字节粒度
cute::make_tma_copy_desc(
gK, sK, tma_k_desc,
Int{});
```
利用Hopper的Tensor Memory Accelerator (TMA) 加速大块数据传输

6. **软流水线优化**
```csrc/flash_fwd_mla_kernel.h
// 修改前
#pragma unroll 1
for (int n_block = ...) {
// 计算阶段
}

// 修改后
#pragma unroll 2
for (int n_block = ...) {
// 阶段1: 加载K/V
// 阶段2: 计算GEMM
// 阶段3: 存储O
}
```
引入软件流水线策略,增加指令级并行

7. **混合精度计算优化**
```csrc/flash_fwd_mla_kernel.h
// 修改前
Tensor rP = flash::convert_type(tSrS);

// 修改后
#if __CUDA_ARCH__ >= 900
Tensor rP = __builtin_convertvector(tSrS, Element); // 使用Hopper原生类型转换
#endif
```
利用Hopper硬件加速的混合精度转换指令

8. **动态资源分配**
```csrc/flash_fwd_mla_kernel.h
// 修改前
constexpr int kNThreads = 256;

// 修改后
constexpr int kNThreads = [](){
if (kHeadDim >= 512) return 256;
else if (kHeadDim >= 256) return 128;
return 64;
}();
```
根据头维度动态调整线程数,优化资源利用率

9. **指令调度优化**
```csrc/flash_fwd_mla_kernel.h
// 修改前
cute::gemm(tiled_mma, tSrQ, tSrK, tSrS);

// 修改后
asm volatile("mma.sync.aligned.m16n8k32.bf16.bf16.bf16 {%0,%1}, {%2}, {%3}, {%4,%5};\n"
: "=r"(tSrS[0]), "=r"(tSrS[1])
: "r"(tSrQ[0]), "r"(tSrK[0]), "r"(tSrS[0]), "r"(tSrS[1]));
```
手工调优关键矩阵乘法的PTX指令调度

10. **缓存提示指令**
```csrc/flash_fwd_mla_kernel.h
// 新增缓存提示
cute::prefetch_global_to_l2(params.q_ptr);
cute::prefetch_global_to_l2(params.k_ptr);
```
利用Hopper L2缓存控制指令优化数据局部性

**性能预期**:
通过上述优化组合,预计可在以下方面提升:
- 内存带宽利用率提升15-20%,达到3500+ GB/s
- 计算吞吐量提升10-15%,达到650+ TFLOPS
- 指令发射效率提升20-30%
- 寄存器压力降低15-20%

建议使用Nsight Compute进行迭代验证,重点关注:
1. stall_long_sb 指标(内存等待)
2. smsp__sass_inst_executed 指令分布
3. l1tex__t_sectors_pipe_lsu_mem_global_op_ld 全局加载效率

Contributor guide

No contributing guide indexed for this repository

Research direction

Start by reading csrc/flash_fwd_mla_kernel.h and profiling the current kernel with Nsight Compute. Compare the proposed changes against stall_long_sb, smsp__sass_inst_executed, and l1tex__t_sectors_pipe_lsu_mem_global_op_ld; the work is done only when validated improvements are demonstrated without breaking the kernel.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
machine-learning, performance
Issue type
Refactor
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.