【训练营】FlashAttention接入 @simon_chou#128

Open

Simon-CHOU wants to merge 33 commits intoInfiniTensor:masterfrom

Simon-CHOU:simon_chou-20260316

Commits on Mar 10, 2026

feat(config): support flash attention configuration
Simon Chou
committed
feat(nn): add ScaledDotProductAttention functional interface
Simon Chou
committed
feat(cli): add --flash command line argument and update model config
Simon Chou
committed
feat(model): integrate flash attention branch in forward pass
Simon Chou
committed
docs: add audit/problems.log.md to track issues
Simon Chou
committed
feat(autograd): implement ScaledDotProductAttention function skeleton
Simon Chou
committed
docs: move tasking plan to audit directory
Simon Chou
committed
feat(kernel): add flash attention forward kernel stub and integrate with autograd
Simon Chou
committed
feat(kernel): implement naive flash attention forward kernel (Br=32, Bc=32, D=64)
Simon Chou
committed
feat(flash-attn): add backward kernel path and validate gpt2 training
Simon Chou
committed

Commits on Mar 12, 2026

test: update benchmark script for stability and llama3 support
Simon Chou
committed
feat: support flash attention flag in llama3 from_llmc and enforce contiguous inputs
Simon Chou
committed
docs: update performance report and problem log for story 4
Simon Chou
committed
fix: resolve LLaMA-3 FlashAttention accuracy issue by correcting input layout and kernel dimensions
Simon Chou
committed
docs: update problem log and mark story 5 as completed
Simon Chou
committed
feat: optimize FlashAttention kernel with tiling and float4, support longer sequence
Simon Chou
committed

Commits on Mar 15, 2026

fix(kernels): resolve NaN issue in FlashAttention WMMA kernel
Simon Chou
committed
test(flash): update test case input range and logic
Simon Chou
committed
docs(audit): update problems log and tasking plan
Simon Chou
committed
perf(kernels): optimize Backward Pass by reducing global atomicAdd
Simon Chou
committed
perf(kernels): optimize Backward Pass with Tiled Accumulation
Simon Chou
committed

Commits on Mar 16, 2026

docs: migrate cuda-report to project root
Simon Chou
committed
fix(kernels): enhance flash attention precision with double accumulator and strict nvcc flags
Simon Chou
committed
test(kernels): add fp32 reference kernel and precision validation test
Simon Chou
committed
fix(runtime): improve cuda error handling and model initialization stability
Simon Chou
committed
feat(examples): add llama3 reproduction scripts and dockerfile
Simon Chou
committed
fix(kernels): align FlashAttention precision with baseline
Simon Chou
committed
docs(audit): log LLaMA-3 performance and precision issues
Simon Chou
committed
docs(report): update judge results to v0.5.0
Simon Chou
committed
chore(git): update gitignore to exclude build artifacts and track report logs
Simon Chou
committed
docs(report): add v0.5.0 performance go/no-go decision document and supporting logs
Simon Chou
committed
docs(report): add tl;dr reproduction instructions for reviewer
Simon Chou
committed
docs(report): embed GPT-2 loss alignment chart
Simon Chou
committed