From a8a3cccf438641802ac6adbbf4935952ae99d6fc Mon Sep 17 00:00:00 2001 From: Kristopher Clark Date: Thu, 23 Jul 2026 09:05:26 -0500 Subject: [PATCH] attn: enable Split-K for mix-batch decode rows Chunked-prefill serving left long-KV decode rows with num_kv_splits=1. Use get_num_splits and allocate the split temporary output so mix-batch decode can parallelize along K like pure decode. Signed-off-by: Kristopher Clark Co-authored-by: Cursor --- csrc/flash_attn/flash_api.cpp | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/csrc/flash_attn/flash_api.cpp b/csrc/flash_attn/flash_api.cpp index a7a4e41779..3ca4b47610 100644 --- a/csrc/flash_attn/flash_api.cpp +++ b/csrc/flash_attn/flash_api.cpp @@ -276,14 +276,28 @@ std::vector mha_varlen_fwd( int effective_seqlen_k = is_local ? std::min(max_seqlen_k, eff_window_left + 1) : max_seqlen_k; - int num_tokens = batch_size; + int num_tokens = q.size(0); int num_heads_q = q.size(1); - int head_dim = q.size(2); int num_heads_kv = k.size(2); int kv_block_size = k.size(1); + int v_head_dim = v.size(-1); - int num_kv_splits = 1; - at::Tensor tmp_out = out; + // Enable Split-K on the decode half of mix-batch. Long KV rows in + // chunked-prefill serving previously forced num_kv_splits=1 and left + // Xe cores idle along K. + int num_kv_splits = num_splits.value_or(get_num_splits( + queue, + batch_size, + num_heads_q, + num_heads_kv, + effective_seqlen_k, + kv_block_size)); + at::Tensor tmp_out = + num_kv_splits == 1 + ? out + : at::empty( + {num_tokens, num_heads_q * num_kv_splits, v_head_dim}, + q.options().device(q.device())); at::Tensor decode_max_logits = at::empty( {num_tokens, num_heads_q, num_kv_splits}, q.options().dtype(at::kFloat).device(q.device()));