[LA] Lightning Attention MTP decode + KVBuffer parallel verify / commit - #97

Merged

fkuner merged 17 commits into

inclusionAI:mainfrom

fkuner:la-decode-kvbuffer

Jul 9, 2026

Commits on Jun 21, 2026

feat: LA decode MTP kernel + tests + benchmark
fankun.fan
committed
57e1729
feat: LA KVBuffer verify + state-update kernels + tests + benchmark
fankun.fan
committed
b73d78d
test: cover odd T for KVBuffer verify/update; fix shuffle SMEM size
fankun.fan
committed
776212d

Commits on Jun 23, 2026

refactor: rename la_update_kvbuffer.py -> la_state_update_kvbuffer.py
fankun.fan
committed
f0179ec
refactor: clean up la_decode_mtp kernel structure
fankun.fan
committed
f65e03c
chore: fix pre-commit issues and inline _la_mtp_ref into test files
fankun.fan
committed
0e11029
chore: add defensive bounds checks for T and V in LA kernels

fankun.fan
authored and
范坤
committed
afad6d4
chore: remove vestigial use_smem_v and use_packed_fma from MMA verify kernel

fankun.fan
authored and
范坤
committed
9e8788b
refactor: collapse ilp_rows branches in la_decode_mtp into generic constexpr loop

fankun.fan
authored and
范坤
committed
d6cde5f
harden LA verify/state-update kernels
fankun.fan
committed
ffd73b9
fix: pass correct cache key to MMA verify kernel in benchmark

fankun.fan
authored and
范坤
committed
5cd5eed
Reuse benchmark_cuda_fn from utils in LA benchmarks.

fankun.fan
authored and
范坤
committed
2ea7d72
Simplify LA MTP/KVBuffer benchmarks with layered compile-cache helpers.

fankun.fan
authored and
范坤
committed
a773efa

Commits on Jun 24, 2026

perf(la): grid-search tile configs per kernel; drop use_smem_v
fankun.fan
committed
7243e86

Commits on Jul 9, 2026

fix(lightning): enforce fp32 for LA MTP and KVBuffer paths

范坤
authored and
范坤
committed
78288ee
Merge upstream main into la-decode-kvbuffer

范坤
authored and
范坤
committed
0608e50
chore: apply Ruff formatting

范坤
authored and
范坤
committed
9a3aec6