Skip to content

Optimize KDA prepare cu_seqlens scan#13

Open
lyppg wants to merge 1 commit into
MoonshotAI:masterfrom
bytedance-iaas:jxp/kda_tune_1
Open

Optimize KDA prepare cu_seqlens scan#13
lyppg wants to merge 1 commit into
MoonshotAI:masterfrom
bytedance-iaas:jxp/kda_tune_1

Conversation

@lyppg

@lyppg lyppg commented Jul 24, 2026

Copy link
Copy Markdown

In current varlen mode, every CTA of _flash_kda_fwd_prepare maps its global_tile_idx to a sequence index by linearly scanning cu_seqlens from the beginning. This PR precomputes a prefix sum of per-sequence tile counts once and have each CTA locate its sequence with an O(log N) binary search.
Result:
Tested on GB200, under varlen, N x 1024 tokens, H=96, D=128, bf16.

N tokens before (ms) after (ms) latency reduction speedup
8 8,192 0.6884 0.6857 0.4% 1.00x
16 16,384 1.3377 1.2980 3.0% 1.03x
32 32,768 2.6996 2.4867 7.9% 1.09x
64 65,536 6.0072 4.8818 18.7% 1.23x
128 131,072 15.2356 9.8047 35.6% 1.55x
256 262,144 44.0999 19.8455 55.0% 2.22x

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant