[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
feat: LA KVBuffer verify + state-update kernels + tests + benchmark
fankun.fan
committed
test: cover odd T for KVBuffer verify/update; fix shuffle SMEM size
fankun.fan
committed

Commits on Jun 23, 2026

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

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

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

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

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

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

fankun.fan
authored and
范坤
committed

Commits on Jun 24, 2026

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

Commits on Jul 9, 2026

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

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

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

范坤
authored and
范坤
committed