feat(kda/sm90): FlashKDA prefill port + intracard-CP prefill - #103

Closed

cherhh wants to merge 45 commits into

inclusionAI:mainfrom

cherhh:flashkda-cp

Commits on Jul 16, 2026

[KDA] port the SM90 fused prefill to a CuTeDSL two-kernel pipeline
cheheng.ch
committed
dabb10b
test sm90 prefill accuracy against FLA triton
cheheng.ch
committed
2826afa
sm90: reject GVA inputs until native support lands
cheheng.ch
committed
4c9e569
reorganize sm90 tests and benches
cheheng.ch
committed
7395d4b
rename hopper_fused_fwd.py -> hopper_prefill.py
cheheng.ch
committed
7e102d9
clean up the sm90 wrapper: import cycle, validation, fp32-only state
cheheng.ch
committed
bb096cb
consolidate sm90 test files
cheheng.ch
committed
4d23c3d
[KDA] add intra-card CP for the sm90 prefill
cheheng.ch
committed
a1d3112
bench: iqr-mean timing + sm90 cp vs fla comparison
cheheng.ch
committed
7c74e3f
fill the whole SM array when splitting long sequences
cheheng.ch
committed
371277a
remove unreachable asserts in internal launchers
cheheng.ch
committed
c9f263a
handle varlen beta in-kernel via K1-emitted ws_beta
cheheng.ch
committed
0213873
tidy cp comments and pad sentinels
cheheng.ch
committed
7abe116
rewrite cula/kda README for users
cheheng.ch
committed
8b5e7c6
test cp against fla ground truth + determinism
cheheng.ch
committed
29f56e0
simplify arch-context and varlen metadata helpers
cheheng.ch
committed
649edcc
fix ws_beta sizing in intracard cp
cheheng.ch
committed
0cd2aef
split the cp driver into plan + pipeline steps
cheheng.ch
committed
f703e83
simplify cp tests, add cp-on/off speedup bench
cheheng.ch
committed
ad5fa3f
assert varlen args in sm100 bwd_wy_dqkg_fused
cheheng.ch
committed
bc1d714
fix varlen jit compile storm (config-keyed cache)
cheheng.ch
committed
0e616cf
add sm90 benchmark numbers and cp debug notes
cheheng.ch
committed
dc7f5d3
drop from-future-annotations in sm90 kernel files (cutlass 4.4.2 Constexpr)
cheheng.ch
committed
52ecfd1
fix K2 init/final state fakes sharing one dynamic sym
cheheng.ch
committed
8644437
launch K1/K2 through tvm-ffi instead of JitExecutor
cheheng.ch
committed
21c0138
pool sm90 workspaces in a grow-only arena
cheheng.ch
committed
4cf9803
cut cp driver host overhead (968 -> 288 us)
cheheng.ch
committed
f0de65c
pre_scan: issue kd@M ahead of the S-chain
cheheng.ch
committed
96aab00
cost-model engage decision for cp auto mode
cheheng.ch
committed
a01e751
same future-annotations fix for the sm100 cp merge kernel
cheheng.ch
committed
89d95cf
slim the prefill wrapper hot path
cheheng.ch
committed
f387dc6
K1 reads beta straight from its packed [T,H] layout
cheheng.ch
committed
1d1a4d6
duration-balanced segment planning for ragged batches
cheheng.ch
committed
44886f3
launch longest segments first
cheheng.ch
committed
f8962f3
test ragged auto-plan cp + workspace arena reuse
cheheng.ch
committed
cb55304
trim redundant comments and stale docs
cheheng.ch
committed
dc9946c
rework cp dispatch: fold the policy layer into the planner
cheheng.ch
committed
0513762
unify copyright headers
cheheng.ch
committed
594a884
slim cp dispatch comments
cheheng.ch
committed
3208b28
style: apply ruff formatting
cheheng.ch
committed
f7cebfa
planner polish: clearer names, worked example, drop the env-warning shim
cheheng.ch
committed
3fd6837
reuse the cpu cu_seqlens copy in sm100 cp dispatch
cheheng.ch
committed
eead625
[KDA] keep the CUDA C++ fused prefill as a separate backend
cheheng.ch
committed
74ba7d4
rename the cutedsl backend export to flashkda_prefill
cheheng.ch
committed
29c04e0

Commits on Jul 17, 2026

drop dead workspace helpers and decorative comments
cheheng.ch
committed
da7534a