fix(fmha): support >4GB KV cache in batch prefill via runtime dispatch#2893

Merged

valarLip merged 12 commits into

mainfrom

jeff-huang/batch-prefill-fix-kvcache-4gb-overflow

Apr 24, 2026

Commits on Apr 24, 2026

test: expand test_batch_prefill_large_kvcache for >4GB KV cache overflow
Jeff-Huang
committed
test: add GPU sync after CK kernel in large_kvcache test
Jeff-Huang
committed
feat(fmha): runtime dispatch for >4GB KV cache in batch prefill
Jeff-Huang
committed
fix: remove unused k_vector_size variable in large_kvcache test
Jeff-Huang
committed
fix(mha): improve batch_prefill TORCH_CHECK error message for >4GB KV cache
Jeff-Huang
committed
test: update scatter_pages comment in large_kvcache test
Jeff-Huang
committed
fix(mha): widen batch_prefill 64-bit threshold to total KV bytes
Jeff-Huang
committed
docs(mha): unify >2GB wording in batch_prefill error and test
Jeff-Huang
committed
refactor(mha): drop wrapper-side use_64bit_load; let CK dispatcher decide
Jeff-Huang
committed
chore(mha): bump CK + update wrapper wording for kUseGlobalLoad rename
Jeff-Huang
committed
test(fmha): parametrize test_batch_prefill_large_kvcache over batch_size {1, 4}
Jeff-Huang
committed
Merge branch 'main' into jeff-huang/batch-prefill-fix-kvcache-4gb-overflow
gyohuangxin
authored