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

Open

fkuner wants to merge 11 commits into

inclusionAI:mainfrom

fkuner:la-decode-kvbuffer

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
refactor: rename la_update_kvbuffer.py -> la_state_update_kvbuffer.py

authored and
范坤
committed
refactor: clean up la_decode_mtp kernel structure

authored and
范坤
committed
chore: fix pre-commit issues and inline _la_mtp_ref into test files

authored and
范坤
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

范坤
authored and
范坤
committed

Commits on Jun 22, 2026

fix: pass correct cache key to MMA verify kernel in benchmark

authored and
范坤
committed