FlashAttention KV Cache Decode — MetaX C500 算子优化工作区

任务概述

实现 FlashAttention GQA KV cache decode 前向算子，超越 baseline。

目标硬件: MetaX C500 (MXMACA) — 104 AP, Wave64, 64KB SMEM/AP, 32 banks
实现语言: TileLang 0.1.10 (MACA backend)
工具链: MACA 3.7.1.5, mxcc, mcProfiler, mcTracer
Baseline: flash_attn_with_kvcache(num_splits=0)

目录结构

/data/opt-flashattn-decode/
├── README.md / SKILL.md / HINTS.md / ITERATIONS.md
├── docs/                          # 工作流 / 版本 / manifest / 合同 / 计划
├── solution/
│   ├── tilelang/                  # TileLang 实现(run_kernel.py = opt_014)
│   └── cuda/                      # CUDA C++ 实现(待开发)
├── bench/                         # 正确性 / benchmark / 编译 / 矩阵 / profile 驱动
├── knowledge/                     # 算子分析 / 模式 / 历史 / hardware-kb(含 MXMACA 内建函数)
├── results/
│   ├── tilelang/                  # TileLang 结果(bench csv + profiling)
│   └── cuda/                      # C++ 结果(待开发)
├── scripts/                       # env / run_bench / run_all / run_profile / summarize_trace
└── workflows/prompts/             # Agent prompt 模板

Baseline 性能

Case	Batch	SeqK	KVhd	GQA	Baseline(ms)	BW(GB/s)	类型
1	1	1	4	8	0.036	0.5	edge
2	4	2	8	4	0.035	2.6	edge
3	16	17	4	8	0.042	14.7	edge
4	64	64	8	4	0.050	236.5	perf
5	16	141	4	8	0.043	75.9	perf
6	16	362	8	4	0.048	322.5	perf
7	64	2048	8	4	0.302	1133.2	perf
8	16	4096	4	8	0.105	829.0	perf
9	32	4096	8	4	0.316	1077.7	perf
10	1	8192	4	8	0.062	272.8	perf
11	16	12251	4	8	0.244	1059.6	perf
12	8	32768	8	4	0.595	1192.9	perf
13	1	58966	8	4	0.212	1137.3	perf
14	1	61519	4	8	0.162	776.2	perf

评测管线

export PYTHONPATH=/opt/tilelang-metax-v0.1.10:/data/TileOPs-Metax:$PYTHONPATH

# Benchmark (AKO4ALL 风格: 冷缓存 + 结构化输出)
CANDIDATE=solution/tilelang/run_kernel.py CASES=10 bash scripts/run_bench.sh

# 正确性+性能对比 (KDA 风格: candidates.jsonl 追踪)
python3 bench/test_tilelang_candidate.py solution/tilelang/run_kernel.py

# 机器可解析输出 (用于 Agent 程序化解析)
python3 bench_run.py --case 10 --cold-cache --machine
# → COMPILED:True CORRECT:True BACKEND:tilelang RUNTIME:0.0623 REF_RUNTIME:0.0618 SPEEDUP:0.9912x

Profiling 管线

# mcProfiler (kernel 级别计数器)
mcProfiler --config mcProfiler.json -- python3 bench_run.py --case 10

# mcTracer (timeline trace)
mcTracer --output results/tilelang/profiling/mcTracer/trace.json -- python3 bench_run.py --case 10

# 诊断检查 (六维分析 + 诊断模式)
# 参见: knowledge/profiling/diagnosis_checklist.md

优化流程

AKO4ALL 风格优化循环

Profile → Edit → Benchmark → Log → Commit → Repeat

每次迭代:

Signal (快): --machine --cold-cache 仅看 runtime 排名
Verdict (慢): 完整对比 + correctness + candidates.jsonl
Log: ITERATIONS.md (Hypothesis → Changes → Bench → Analysis → Next)
Commit: git commit -m "[iter N] <description>"
Trajectory: 自动备份 solution/ + results 到 trajectory/

停滞处理 (3次无改进):

重新 profiling (mcProfiler)
检查诊断模式 (diagnosis_checklist.md A-M)
搜索最优已知技术

候选者追踪 (candidates.jsonl)

{"id":"candidate-20260815_123456-case10","timestamp":"...",
 "technique":"split-kv-v1","validation":"pass","runtime_ms":0.045,
 "ref_runtime_ms":0.062,"speedup":1.378,"config":"block_H=32,block_N=32,..."}

KDA 六维分析框架

维度	内容	工具
1. Occupancy & Launch Geometry	CTA 数、waves/AP	mcProfiler
2. Tail Effect	变长序列负载不均衡	mcTracer timeline
3. Stall Breakdown	long_scoreboard / barrier	mcProfiler
4. Tensor Core Utilization	GEMM 活跃度	mcProfiler CE
5. SM Timeline	kernel 执行形状	mcTracer
6. Memory Bandwidth	实际 vs 理论 BW	bench_run.py

关键优化方向

Case 10, 14 (batch=1 KV4): 273-776 GB/s — 小网格 (诊断模式A)
Case 13 (batch=1 KV8): 1137 GB/s — 已高，可进一步调优
Split-KV — 增加 CTA 数铺满 104 AP
GQA Grouped KV Read — SMEM 复用 K/V
SMEM Budget — block_H=32, block_N=32 (36KB, 2 blocks/AP)
Bank Conflict Pad — headdim=128 → stride=132
Cold-Cache Timing — L2 冲刷模拟真实场景

参考仓库

AKO4ALL — Agentic Kernel Optimization Loop
KDA — Kernel Design Agents + KernelWiki
DocFlowAI — MetaX MACA 文档生态

来自队友的技能知识库 (metax-skills-20260812)

本工作区已将队友整理的技能包 /Users/chenxi/Desktop/metax-skills-20260812/ 完整整合。

整合内容

模块	工作区位置	用途
C500 硬件架构	`knowledge/hardware-kb/c500/`	AP 104, Wave64, 64KB SMEM, L2 8MB
MACA 内建函数	`knowledge/hardware-kb/maca/builtins.md`	MMA, shuffle, memory builtins
mxcc 编译器约束	`knowledge/hardware-kb/maca/mxcc_gqa.md`	GQA 共享 KV 编译器限制
算子数学分析 (L0)	`knowledge/operator-analysis.md`	GQA paged decode 计算图与理论极限
可移植性检查清单	`knowledge/portability-checklist.md`	CUDA→MACA 迁移检查项
三层优化模式库	`knowledge/optimization-patterns/`	按 layers (L0-L3) 组织的变换模式
Profiling 诊断框架	`knowledge/profiling/`	mcProfiler + 六维分析 + 13 诊断模式
环境变量模板	`scripts/env.sh`	PYTHONPATH, MACA_PATH 一键配置
项目 Skills	`.claude/skills/`	优化/Profiling/审查三个可调用指令
知识库索引	`knowledge/README.md`	按需检索指引

层级架构 (L0-L3)

层	内容	在本项目中的应用
L0 Math	数学等价性、数值分析	online softmax 单遍等价性、GQA 计算强度
L1 Graph	图融合、折叠、重排	QK→softmax→PV 融合、split-KV 子图拆分
L2 Algorithm	算法结构	GQA group dispatch、host schedule materialization
L3 Kernel	分块、线程、SMEM、同步	flash attention tiling、SMEM bank conflict pad
Filters	Profiling 瓶颈分类	小网格、尾效应、延迟绑定、bank conflict
Paradigm	重写范式	持久化 kernel、工作窃取、warp specialization