diff --git a/include/sgl_flash_kernel_ops.h b/include/sgl_flash_kernel_ops.h index 44b9e61a4..c3a28c445 100644 --- a/include/sgl_flash_kernel_ops.h +++ b/include/sgl_flash_kernel_ops.h @@ -74,7 +74,10 @@ std::tuple mha_fwd( int num_kv_splits, std::optional pack_gqa_, int const sm_margin, - std::optional& out_); + std::optional& out_, + std::optional& k_new_, + std::optional& v_new_, + std::optional& cu_seqlens_k_new_); void flash_mla_decode( torch::Tensor& out, diff --git a/python/sgl_kernel/flash_attn.py b/python/sgl_kernel/flash_attn.py index 78740c5c8..930d4be32 100644 --- a/python/sgl_kernel/flash_attn.py +++ b/python/sgl_kernel/flash_attn.py @@ -264,6 +264,14 @@ def flash_attn_with_kvcache( if cache_seqlens is not None: assert cache_seqlens.size(0) + 1 == cu_seqlens_q.size(0) cu_seqlens_k = cache_seqlens + has_new_kv = k is not None or v is not None or cu_seqlens_k_new is not None + native_max_seqlen_k = 1 + if has_new_kv: + native_max_seqlen_k = max_seqlen_k + if native_max_seqlen_k is None or native_max_seqlen_k == 0: + native_max_seqlen_k = ( + k.shape[1] if k is not None and k.dim() == 4 else (max_seqlen_q or 1) + ) out, softmax_lse, *rest = torch.ops.sgl_kernel.fwd.default( q, k_cache, @@ -272,7 +280,7 @@ def flash_attn_with_kvcache( cu_seqlens_q, cu_seqlens_k, max_seqlen_q, - 1, + native_max_seqlen_k, page_table, cache_batch_idx, cache_leftpad, @@ -294,6 +302,9 @@ def flash_attn_with_kvcache( pack_gqa, sm_margin, out, + k, + v, + cu_seqlens_k_new, ) return (out, softmax_lse, *rest) if return_softmax_lse else out diff --git a/src/FMHAPrefillXe20.cmake b/src/FMHAPrefillXe20.cmake index 9e96270dc..ed0315648 100644 --- a/src/FMHAPrefillXe20.cmake +++ b/src/FMHAPrefillXe20.cmake @@ -31,19 +31,40 @@ set(FMHA_PREFILL_NUM_SG_64 8) set(FMHA_PREFILL_TILED_Q_96 128) set(FMHA_PREFILL_TILED_KV_96 64) -set(FMHA_PREFILL_NUM_SG_96 8) +set(FMHA_PREFILL_NUM_SG_96 16) -set(FMHA_PREFILL_TILED_Q_128 256) -set(FMHA_PREFILL_TILED_KV_128 32) +set(FMHA_PREFILL_TILED_Q_128 128) +set(FMHA_PREFILL_TILED_KV_128 64) set(FMHA_PREFILL_NUM_SG_128 16) +option(FMHA_PREFILL_HD128_LARGE_TILE "Enable q256/k32 paged head_dim=128 path for model-sized Q lengths" ON) +set(FMHA_PREFILL_HD128_LARGE_TILE_MIN_Q 256) +set(FMHA_PREFILL_HD128_LARGE_TILE_Q 256) +set(FMHA_PREFILL_HD128_LARGE_TILE_KV 32) +set(FMHA_PREFILL_HD128_LARGE_NUM_SG 16) -set(FMHA_PREFILL_TILED_Q_192 256) +set(FMHA_PREFILL_TILED_Q_192 128) set(FMHA_PREFILL_TILED_KV_192 64) -set(FMHA_PREFILL_NUM_SG_192 32) +set(FMHA_PREFILL_NUM_SG_192 16) -set(FMHA_PREFILL_TILED_Q_256 256) +set(FMHA_PREFILL_TILED_Q_256 128) set(FMHA_PREFILL_TILED_KV_256 64) -set(FMHA_PREFILL_NUM_SG_256 32) +set(FMHA_PREFILL_NUM_SG_256 16) +option(FMHA_PREFILL_HD256_LARGE_TILE "Enable q256 paged head_dim=256 path for large AppendKV prefill" ON) +set(FMHA_PREFILL_HD256_LARGE_TILE_MIN_Q 256) +set(FMHA_PREFILL_HD256_LARGE_TILE_Q 256) +set(FMHA_PREFILL_HD256_LARGE_TILE_KV 64) +set(FMHA_PREFILL_HD256_LARGE_NUM_SG 32) + +set(FMHA_PREFILL_APPENDKV_Q256_HEAD_DIMS 64 192 256) +set(FMHA_PREFILL_APPENDKV_Q256_Q_64 256) +set(FMHA_PREFILL_APPENDKV_Q256_KV_64 64) +set(FMHA_PREFILL_APPENDKV_Q256_NUM_SG_64 16) +set(FMHA_PREFILL_APPENDKV_Q256_Q_192 256) +set(FMHA_PREFILL_APPENDKV_Q256_KV_192 64) +set(FMHA_PREFILL_APPENDKV_Q256_NUM_SG_192 32) +set(FMHA_PREFILL_APPENDKV_Q256_Q_256 256) +set(FMHA_PREFILL_APPENDKV_Q256_KV_256 64) +set(FMHA_PREFILL_APPENDKV_Q256_NUM_SG_256 32) set(FMHA_PREFILL_TILED_Q_512 256) set(FMHA_PREFILL_TILED_KV_512 64) @@ -66,9 +87,19 @@ set(FMHA_PREFILL_TILED_Q_NP_80 256) set(FMHA_PREFILL_TILED_KV_NP_80 64) set(FMHA_PREFILL_NUM_SG_NP_80 16) -set(FMHA_PREFILL_TILED_Q_NP_96 256) +set(FMHA_PREFILL_TILED_Q_NP_96 128) set(FMHA_PREFILL_TILED_KV_NP_96 64) set(FMHA_PREFILL_NUM_SG_NP_96 16) +option(FMHA_PREFILL_HD96_NP_SMALL_TILE "Enable q32 non-paged head_dim=96 path for short queries" ON) +set(FMHA_PREFILL_HD96_NP_SMALL_TILE_MAX_Q 32) +set(FMHA_PREFILL_HD96_NP_SMALL_TILE_Q 32) +set(FMHA_PREFILL_HD96_NP_SMALL_TILE_KV 64) +set(FMHA_PREFILL_HD96_NP_SMALL_NUM_SG 4) +option(FMHA_PREFILL_HD96_NP_LARGE_TILE "Enable q256 non-paged head_dim=96 path for long queries" ON) +set(FMHA_PREFILL_HD96_NP_LARGE_TILE_MIN_Q 257) +set(FMHA_PREFILL_HD96_NP_LARGE_TILE_Q 256) +set(FMHA_PREFILL_HD96_NP_LARGE_TILE_KV 64) +set(FMHA_PREFILL_HD96_NP_LARGE_NUM_SG 16) set(FMHA_PREFILL_TILED_Q_NP_128 256) set(FMHA_PREFILL_TILED_KV_NP_128 32) @@ -96,6 +127,39 @@ foreach(HEAD_DIM ${FMHA_PREFILL_PAGED_HEAD_DIMS}) set(TILED_OUT ${FMHA_PREFILL_TILED_OUT_${HEAD_DIM}}) endif() + if(HEAD_DIM STREQUAL "128" AND FMHA_PREFILL_HD128_LARGE_TILE) + set(HD128_PAGED_LARGE_TILE 1) + else() + set(HD128_PAGED_LARGE_TILE 0) + endif() + set(HD128_PAGED_LARGE_TILE_MIN_Q ${FMHA_PREFILL_HD128_LARGE_TILE_MIN_Q}) + set(HD128_PAGED_LARGE_TILE_Q ${FMHA_PREFILL_HD128_LARGE_TILE_Q}) + set(HD128_PAGED_LARGE_TILE_KV ${FMHA_PREFILL_HD128_LARGE_TILE_KV}) + set(HD128_PAGED_LARGE_NUM_SG ${FMHA_PREFILL_HD128_LARGE_NUM_SG}) + + if(HEAD_DIM STREQUAL "256" AND FMHA_PREFILL_HD256_LARGE_TILE) + set(HD256_PAGED_LARGE_TILE 1) + else() + set(HD256_PAGED_LARGE_TILE 0) + endif() + set(HD256_PAGED_LARGE_TILE_MIN_Q ${FMHA_PREFILL_HD256_LARGE_TILE_MIN_Q}) + set(HD256_PAGED_LARGE_TILE_Q ${FMHA_PREFILL_HD256_LARGE_TILE_Q}) + set(HD256_PAGED_LARGE_TILE_KV ${FMHA_PREFILL_HD256_LARGE_TILE_KV}) + set(HD256_PAGED_LARGE_NUM_SG ${FMHA_PREFILL_HD256_LARGE_NUM_SG}) + + list(FIND FMHA_PREFILL_APPENDKV_Q256_HEAD_DIMS ${HEAD_DIM} APPENDKV_Q256_INDEX) + if(APPENDKV_Q256_INDEX GREATER -1) + set(APPENDKV_Q256_TILE 1) + set(APPENDKV_Q256_Q ${FMHA_PREFILL_APPENDKV_Q256_Q_${HEAD_DIM}}) + set(APPENDKV_Q256_KV ${FMHA_PREFILL_APPENDKV_Q256_KV_${HEAD_DIM}}) + set(APPENDKV_Q256_NUM_SG ${FMHA_PREFILL_APPENDKV_Q256_NUM_SG_${HEAD_DIM}}) + else() + set(APPENDKV_Q256_TILE 0) + set(APPENDKV_Q256_Q 256) + set(APPENDKV_Q256_KV 64) + set(APPENDKV_Q256_NUM_SG 16) + endif() + set(GENERATED_FILE "${CMAKE_CURRENT_BINARY_DIR}/sycl/xe_fmha_fwd_prefill_paged_kernel_${HEAD_DIM}.cpp") configure_file(${FMHA_PREFILL_TEMPLATE} ${GENERATED_FILE} @ONLY) @@ -116,6 +180,26 @@ foreach(HEAD_DIM ${FMHA_PREFILL_NP_HEAD_DIMS}) message(FATAL_ERROR "Missing non-paged tile params for prefill HEAD_DIM=${HEAD_DIM}") endif() + if(HEAD_DIM STREQUAL "96" AND FMHA_PREFILL_HD96_NP_SMALL_TILE) + set(HD96_NP_SMALL_TILE 1) + else() + set(HD96_NP_SMALL_TILE 0) + endif() + set(HD96_NP_SMALL_TILE_MAX_Q ${FMHA_PREFILL_HD96_NP_SMALL_TILE_MAX_Q}) + set(HD96_NP_SMALL_TILE_Q ${FMHA_PREFILL_HD96_NP_SMALL_TILE_Q}) + set(HD96_NP_SMALL_TILE_KV ${FMHA_PREFILL_HD96_NP_SMALL_TILE_KV}) + set(HD96_NP_SMALL_NUM_SG ${FMHA_PREFILL_HD96_NP_SMALL_NUM_SG}) + + if(HEAD_DIM STREQUAL "96" AND FMHA_PREFILL_HD96_NP_LARGE_TILE) + set(HD96_NP_LARGE_TILE 1) + else() + set(HD96_NP_LARGE_TILE 0) + endif() + set(HD96_NP_LARGE_TILE_MIN_Q ${FMHA_PREFILL_HD96_NP_LARGE_TILE_MIN_Q}) + set(HD96_NP_LARGE_TILE_Q ${FMHA_PREFILL_HD96_NP_LARGE_TILE_Q}) + set(HD96_NP_LARGE_TILE_KV ${FMHA_PREFILL_HD96_NP_LARGE_TILE_KV}) + set(HD96_NP_LARGE_NUM_SG ${FMHA_PREFILL_HD96_NP_LARGE_NUM_SG}) + set(GENERATED_NP_FILE "${CMAKE_CURRENT_BINARY_DIR}/sycl/xe_fmha_fwd_prefill_nopage_kernel_${HEAD_DIM}.cpp") configure_file(${FMHA_PREFILL_NOPAGE_TEMPLATE} ${GENERATED_NP_FILE} @ONLY) diff --git a/src/sycl/flash_attention.cpp b/src/sycl/flash_attention.cpp index e45e238c6..e0f929c66 100644 --- a/src/sycl/flash_attention.cpp +++ b/src/sycl/flash_attention.cpp @@ -183,8 +183,9 @@ std::vector mha_fwd_nopage( params.cu_seqlens_q = cu_seqlens_q.data_ptr(); params.cu_seqlens_k = cu_seqlens_k.data_ptr(); - // No "new" KV: the whole sequence lives in the contiguous cache buffer, so the - // decode kernel reads everything from the K/V cache pointers (knew = 0). + // No "new" KV: the whole sequence lives in the contiguous cache buffer. + params.knew_ptr = nullptr; + params.vnew_ptr = nullptr; params.cu_seqlens_knew = nullptr; params.seqlen_knew = 0; params.total_knew = 0; @@ -213,6 +214,8 @@ std::vector mha_fwd_nopage( } params.softcap = softcap; params.p_dropout = 1.f; + params.is_e4m3 = false; + params.is_e5m2 = false; // Decode never needs a causal mask (each selected batch has seqlen_q <= 1, so // a single query attends to the full cache); sliding-window/local masking is @@ -517,6 +520,11 @@ std::vector mha_fwd( params.cu_seqlens_q = cu_seqlens_q.data_ptr(); params.cu_seqlens_k = cu_seqlens_k.data_ptr(); + params.knew_ptr = nullptr; + params.vnew_ptr = nullptr; + params.cu_seqlens_knew = nullptr; + params.seqlen_knew = 0; + params.total_knew = 0; params.num_kv_splits = num_kv_splits; // Softmax sum @@ -772,6 +780,12 @@ std::vector mha_fwd_nopage( params.cu_seqlens_q = cu_seqlens_q.data_ptr(); params.cu_seqlens_k = cu_seqlens_k.data_ptr(); + params.k_new_ptr = nullptr; + params.v_new_ptr = nullptr; + params.cu_seqlens_kvnew = nullptr; + params.kv_cache_seqlens = nullptr; + params.seqlen_kvnew = 0; + params.total_kvnew = 0; params.softmax_lse_ptr = softmax_lse.data_ptr(); @@ -780,7 +794,7 @@ std::vector mha_fwd_nopage( params.h_k = num_heads_k; params.q_group_size = 1; params.seqlen_q = seqlen_q; - params.seqlen_k = seqlen_k; + params.seqlen_kvcache = seqlen_k; params.d = head_size; params.d_rounded = head_size_rounded; @@ -805,7 +819,7 @@ std::vector mha_fwd_nopage( params.window_size_left = window_size_left; params.window_size_right = window_size_right; params.total_q = total_q; - params.total_k = total_k; + params.total_kvcache = total_k; params.b_k = batch_size; params.dv = head_size_v; @@ -889,7 +903,10 @@ std::vector mha_fwd( // chunkprefill two-launch path: pre-allocated shared output, and a per-batch // bool mask (length = batch) whose true entries are skipped by the kernel. std::optional out_opt = std::nullopt, - std::optional skip_batch_mask_opt = std::nullopt) { + std::optional skip_batch_mask_opt = std::nullopt, + std::optional k_new_ = std::nullopt, + std::optional v_new_ = std::nullopt, + std::optional cu_seqlens_k_new_ = std::nullopt) { auto q_type = q.scalar_type(); TORCH_CHECK( q_type == at::ScalarType::Half || q_type == at::ScalarType::BFloat16, @@ -915,8 +932,11 @@ std::vector mha_fwd( CHECK_LAST_DIM_CONTIGUOUS_INPUT(k); CHECK_LAST_DIM_CONTIGUOUS_INPUT(v); + bool const has_new_kv = k_new_.has_value() || v_new_.has_value() || cu_seqlens_k_new_.has_value(); + // Non-paged (page_table == nullopt) prefill: contiguous ragged KV cache. if (!page_table.has_value()) { + TORCH_CHECK(!has_new_kv, "AppendKV requires paged KV cache"); return mha_fwd_nopage( q, k, @@ -987,6 +1007,15 @@ std::vector mha_fwd( CHECK_SHAPE(v, num_pages, page_size, num_heads_k, head_size_v); CHECK_SHAPE(page_table.value(), batch_size_k, max_num_pages_per_seq); + at::Tensor page_table_for_kernel = page_table.value(); + if (kv_batch_idx_.has_value()) { + auto kv_batch_idx = kv_batch_idx_.value(); + CHECK_INPUT(kv_batch_idx); + TORCH_CHECK(kv_batch_idx.scalar_type() == torch::kInt32, "kv_batch_idx must have dtype int32"); + CHECK_SHAPE(kv_batch_idx, batch_size); + page_table_for_kernel = page_table_for_kernel.index_select(0, kv_batch_idx.toType(torch::kLong)); + } + if (leftpad_k_.has_value()) { auto leftpad_k = leftpad_k_.value(); TORCH_CHECK(leftpad_k.dtype() == torch::kInt32, "leftpad_k must have dtype int32"); @@ -1047,7 +1076,7 @@ std::vector mha_fwd( params.h_k = num_heads_k; params.q_group_size = 1; params.seqlen_q = seqlen_q; - params.seqlen_k = seqlen_k; + params.seqlen_kvcache = seqlen_k; params.d = head_size; params.d_rounded = head_size_rounded; @@ -1092,15 +1121,62 @@ std::vector mha_fwd( params.window_size_left = window_size_left; params.window_size_right = window_size_right; params.total_q = total_q; - params.total_k = total_k; - params.b_k = batch_size_k; + params.total_kvcache = total_k; + params.b_k = page_table_for_kernel.size(0); params.dv = head_size_v; - params.page_table = page_table.value().data_ptr(); - params.page_table_batch_stride = page_table.value().stride(0); + params.page_table = page_table_for_kernel.data_ptr(); + params.page_table_batch_stride = page_table_for_kernel.stride(0); params.max_num_pages_per_seq = max_num_pages_per_seq; params.page_size = page_size; params.num_pages = num_pages; + params.k_new_ptr = nullptr; + params.v_new_ptr = nullptr; + params.cu_seqlens_kvnew = nullptr; + params.kv_cache_seqlens = nullptr; + params.seqlen_kvnew = 0; + params.total_kvnew = 0; + if (has_new_kv) { + TORCH_CHECK(k_new_.has_value() && v_new_.has_value(), "AppendKV requires both k_new and v_new"); + auto const& k_new = k_new_.value(); + auto const& v_new = v_new_.value(); + CHECK_LAST_DIM_CONTIGUOUS_INPUT(k_new); + CHECK_LAST_DIM_CONTIGUOUS_INPUT(v_new); + TORCH_CHECK(k_new.scalar_type() == k.scalar_type(), "k_new dtype must match KV cache key dtype"); + TORCH_CHECK(v_new.scalar_type() == v.scalar_type(), "v_new dtype must match KV cache value dtype"); + TORCH_CHECK(k_new.dim() == 3 || k_new.dim() == 4, "k_new must be [total_k_new, h_k, d] or [b, s, h_k, d]"); + TORCH_CHECK(v_new.dim() == k_new.dim(), "v_new rank must match k_new rank"); + int total_kvnew = 0; + int seqlen_kvnew = max_seqlen_k > 0 ? max_seqlen_k : max_seqlen_q; + if (k_new.dim() == 3) { + total_kvnew = k_new.size(0); + CHECK_SHAPE(k_new, total_kvnew, num_heads_k, head_size); + CHECK_SHAPE(v_new, total_kvnew, num_heads_k, head_size_v); + TORCH_CHECK( + cu_seqlens_k_new_.has_value() || seqlen_kvnew > 0, + "ragged k_new requires cu_seqlens_k_new or positive max_seqlen_k"); + } else { + TORCH_CHECK(k_new.size(0) == batch_size, "batched k_new first dimension must match batch size"); + int const k_new_seqlen = k_new.size(1); + total_kvnew = batch_size * k_new_seqlen; + seqlen_kvnew = max_seqlen_k > 0 ? max_seqlen_k : k_new_seqlen; + CHECK_SHAPE(k_new, batch_size, k_new_seqlen, num_heads_k, head_size); + CHECK_SHAPE(v_new, batch_size, k_new_seqlen, num_heads_k, head_size_v); + } + if (cu_seqlens_k_new_.has_value()) { + auto const& cu_seqlens_k_new = cu_seqlens_k_new_.value(); + CHECK_INPUT(cu_seqlens_k_new); + TORCH_CHECK(cu_seqlens_k_new.dtype() == torch::kInt32, "cu_seqlens_k_new must have dtype torch.int32"); + CHECK_SHAPE(cu_seqlens_k_new, batch_size + 1); + params.cu_seqlens_kvnew = cu_seqlens_k_new.data_ptr(); + } + params.k_new_ptr = k_new.data_ptr(); + params.v_new_ptr = v_new.data_ptr(); + params.kv_cache_seqlens = cu_seqlens_k.data_ptr(); + params.seqlen_kvnew = seqlen_kvnew; + params.total_kvnew = total_kvnew; + } + if (q_v_.has_value()) { TORCH_CHECK(head_size <= 64, "q_v is only supported for head_size <= 64"); TORCH_CHECK( @@ -1233,7 +1309,10 @@ std::vector mha_fwd( int num_kv_splits, std::optional pack_gqa_, int const sm_margin, - std::optional out_ = std::nullopt) { + std::optional out_ = std::nullopt, + std::optional k_new_ = std::nullopt, + std::optional v_new_ = std::nullopt, + std::optional cu_seqlens_k_new_ = std::nullopt) { // Supports both paged (page_table != None) and non-paged (contiguous ragged // KV, page_table == None) layouts. // ``seqlens_rotary_`` is intentionally not checked here: callers pass it @@ -1244,12 +1323,14 @@ std::vector mha_fwd( !scheduler_metadata_.has_value(), "chunkprefill two-launch path does not yet support q_v / rotary / q_descale / scheduler_metadata."); TORCH_CHECK(cu_seqlens_q.scalar_type() == at::kInt, "cu_seqlens_q must be int32."); + bool const has_new_kv = k_new_.has_value() || v_new_.has_value() || cu_seqlens_k_new_.has_value(); // Pre-allocated out requires paged KV: on the non-paged path zero-KV-length // rows are never written by the kernel, so a caller buffer would retain stale // values on graph replay. SGLang always provides page_table (paged KV cache), // so this check should never fire in practice. TORCH_CHECK( !out_.has_value() || page_table.has_value(), "chunkprefill: out buffer requires page_table (paged KV cache)."); + TORCH_CHECK(!has_new_kv || page_table.has_value(), "chunkprefill: AppendKV requires page_table (paged KV cache)."); int64_t batch_size = cu_seqlens_q.size(0) - 1; TORCH_CHECK(batch_size >= 0, "cu_seqlens_q must have at least 1 element."); @@ -1297,7 +1378,40 @@ std::vector mha_fwd( // out_ buffer) and skips prefill batches. auto out = launch(decode::mha_fwd, std::move(out_), is_prefill)[0]; // Launch 2: prefill writes into the same output and skips decode batches. - launch(prefill::mha_fwd, out, is_prefill.logical_not()); + prefill::mha_fwd( + q, + k, + v, + q_v_, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + page_table, + kv_batch_idx_, + leftpad_k_, + rotary_cos_, + rotary_sin_, + seqlens_rotary_, + q_descale_, + k_descale_, + v_descale_, + softmax_scale_, + sinks_, + is_causal, + window_size_left, + window_size_right, + softcap, + is_rotary_interleaved, + scheduler_metadata_, + num_kv_splits, + pack_gqa_, + sm_margin, + out, + is_prefill.logical_not(), + k_new_, + v_new_, + cu_seqlens_k_new_); // softmax_lse / accum tensors are not stitched here; return empty // placeholders to keep the Python ABI stable. @@ -1336,7 +1450,10 @@ std::tuple mha_fwd( int num_kv_splits, std::optional pack_gqa_, int const sm_margin, - std::optional& out_) { + std::optional& out_, + std::optional& k_new_, + std::optional& v_new_, + std::optional& cu_seqlens_k_new_) { TORCH_CHECK(q.dim() == 3, "query must be in ragged format (total_q, h, d)"); // k and v may be 3D (total_k, h_k, d) for non-paged or 4D (num_pages, page_size, h_k, d) // for paged KV cache; sub-functions validate their own shapes. @@ -1354,6 +1471,7 @@ std::tuple mha_fwd( int const num_heads = q.size(-2); int const num_heads_k = k.size(-2); int64_t batch_size = cu_seqlens_q.size(0) - 1; + bool const has_new_kv = k_new_.has_value() || v_new_.has_value() || cu_seqlens_k_new_.has_value(); // decode / prefill / chunkprefill all take the same leading argument list; // only the trailing parameters differ. Bind the shared arguments once here so @@ -1394,7 +1512,9 @@ std::tuple mha_fwd( std::forward(tail)...)); }; - if (max_seqlen_q == 1) { + if (has_new_kv) { + return dispatch(prefill::mha_fwd, std::nullopt, k_new_, v_new_, cu_seqlens_k_new_); + } else if (max_seqlen_q == 1) { // Pure decode path return dispatch(decode::mha_fwd, std::nullopt); } else if (!page_table.has_value() || batch_size == 1) { @@ -1404,11 +1524,11 @@ std::tuple mha_fwd( // is_prefill.all() — a device reduction + D2H sync that costs more than it saves. // But batch_size == 1 makes it provable from host scalars: // a single sequence with max_seqlen_q > 1 is prefill - return dispatch(prefill::mha_fwd, std::nullopt); + return dispatch(prefill::mha_fwd, std::nullopt, std::nullopt, std::nullopt, std::nullopt); } else { // Chunk prefill path // Paged attn with max_seqlen_q > 1 and batch_size > 1 - return dispatch(chunkprefill::mha_fwd); + return dispatch(chunkprefill::mha_fwd, k_new_, v_new_, cu_seqlens_k_new_); } } #undef SYCL_INTEL_TARGET diff --git a/src/sycl/kernels/flash_attention_v2/collective/xe_fmha_fwd_mainloop.hpp b/src/sycl/kernels/flash_attention_v2/collective/xe_fmha_fwd_mainloop.hpp index 4a18bba0b..8cb06da0a 100644 --- a/src/sycl/kernels/flash_attention_v2/collective/xe_fmha_fwd_mainloop.hpp +++ b/src/sycl/kernels/flash_attention_v2/collective/xe_fmha_fwd_mainloop.hpp @@ -31,12 +31,15 @@ #pragma once +#include + #include "cute/algorithm/functional.hpp" #include "cute/algorithm/gemm.hpp" #include "cute/algorithm/subgroup_algorithms.hpp" #include "cute/atom/mma_atom.hpp" #include "cutlass/cutlass.h" #include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/sycl_vector_types.h" #include "fmha_fusion.hpp" namespace cutlass::fmha { @@ -52,6 +55,13 @@ using namespace cute; ///////////////////////////////////////////////////////////////////////////////////////////////// +enum class AppendKVMode { + kNone, + kStore, + kStoreWide, + kStoreAndDirectLoad, +}; + template < class DispatchPolicy_, bool CausalMask_, @@ -75,7 +85,8 @@ template < // (decode only, seq_len_qo == 1). All packed rows share the single decode // KV position, so per-row masking must use a fixed decode row. Default // false keeps prefill (and non-packed decode) unaffected. - bool PackGQA_ = false> + bool PackGQA_ = false, + AppendKVMode AppendKVMode_ = AppendKVMode::kNone> struct FMHAFwdMainloop { static_assert(cutlass::detail::dependent_false, "Could not find a mainloop specialization."); }; @@ -101,7 +112,8 @@ template < class TiledCopyK_cache_, class TiledCopyV_cache_, bool LocalMask_, - bool PackGQA_> + bool PackGQA_, + AppendKVMode AppendKVMode_> struct FMHAFwdMainloop< XeDefault, CausalMask_, @@ -121,7 +133,8 @@ struct FMHAFwdMainloop< TiledCopyK_cache_, TiledCopyV_cache_, LocalMask_, - PackGQA_> { + PackGQA_, + AppendKVMode_> { // // Type Aliases // @@ -199,6 +212,10 @@ struct FMHAFwdMainloop< // K/V are dequantized (cast to ElementQ and multiplied by the per-tensor // scale) inside the mainloop after the block-2D load. static constexpr bool Fp8KV = is_any_of_v; + static constexpr AppendKVMode AppendMode = AppendKVMode_; + static constexpr bool AppendKV = AppendMode != AppendKVMode::kNone; + static constexpr bool AppendDirectLoad = AppendMode == AppendKVMode::kStoreAndDirectLoad && PagedKV; + static constexpr bool AppendWideStore = AppendMode == AppendKVMode::kStoreWide && PagedKV; // User-facing arguments struct Arguments { @@ -208,6 +225,12 @@ struct FMHAFwdMainloop< int max_num_pages_per_seq = 0; int window_size_left = -1; int window_size_right = -1; + typename TensorK_cache::element_type const* ptr_K_new = nullptr; + typename TensorV_cache::element_type const* ptr_V_new = nullptr; + int const* ptr_cu_seqlens_k_new = nullptr; + int const* ptr_kv_cache_seqlens = nullptr; + int seq_len_kv_new = 0; + int total_k_new = 0; }; // Kernel-facing parameters @@ -233,7 +256,13 @@ struct FMHAFwdMainloop< args.page_size, args.max_num_pages_per_seq, args.window_size_left, - args.window_size_right}; + args.window_size_right, + args.ptr_K_new, + args.ptr_V_new, + args.ptr_cu_seqlens_k_new, + args.ptr_kv_cache_seqlens, + args.seq_len_kv_new, + args.total_k_new}; } CUTLASS_HOST_DEVICE static bool can_implement(Arguments const&) { @@ -254,6 +283,275 @@ struct FMHAFwdMainloop< return params.ptr_page_table[batch_offset + next_page_logical_idx] * tiles_per_page + K % tiles_per_page; } + CUTLASS_DEVICE + static int get_kv_len(int const* cu_seqlens, int seq_len, int total_len, int batch) { + if (total_len <= 0 || (cu_seqlens == nullptr && seq_len <= 0)) { + return 0; + } + if (cu_seqlens != nullptr) { + return cu_seqlens[batch + 1] - cu_seqlens[batch]; + } + return seq_len; + } + + CUTLASS_DEVICE + static int get_kv_new_len(Params const& params, int batch) { + if constexpr (AppendKV) { + if (params.ptr_K_new == nullptr || params.ptr_V_new == nullptr || params.ptr_kv_cache_seqlens == nullptr) { + return 0; + } + return get_kv_len(params.ptr_cu_seqlens_k_new, params.seq_len_kv_new, params.total_k_new, batch); + } else { + (void)params; + (void)batch; + return 0; + } + } + + CUTLASS_DEVICE + int get_kv_new_len(int batch) const { + return get_kv_new_len(params, batch); + } + + template + CUTLASS_DEVICE bool store_kv_new_vectorized( + TensorK_cache2D const& K_cache_2D, + TensorV_cache2D const& V_cache_2D, + int batch, + int kv_head, + int num_heads_kv, + int lane_idx, + int sub_group_id, + int new_len, + int new_begin, + int cache_len_old, + int head_size_qk, + int head_size_vo) const { + if (head_size_qk != head_size_vo || (head_size_qk % kVecElems) != 0) { + return false; + } + + auto& K_dst = const_cast(K_cache_2D); + auto& V_dst = const_cast(V_cache_2D); + // Large appends are bandwidth-bound, so split tokens across SGs; + // small appends keep the old single-SG path to avoid control overhead. + constexpr int kMinMultiSgTokens = 64; + int const active_sg_count = new_len >= kMinMultiSgTokens ? int(SGPerWG::value) : 1; + if (sub_group_id >= active_sg_count) { + return true; + } + int const head_size = head_size_qk; + bool const flatten_token_vectors = new_len >= kMinMultiSgTokens && head_size == 64; + int const vecs_per_token = head_size / kVecElems; + int const worker = sub_group_id * intel::sg_size + lane_idx; + int const worker_count = active_sg_count * intel::sg_size; + bool single_dst_page = true; + int single_row_base = cache_len_old; + if constexpr (PagedKV) { + int const first_dst_page = cache_len_old / params.page_size; + int const last_dst_page = (cache_len_old + new_len - 1) / params.page_size; + single_dst_page = first_dst_page == last_dst_page; + if (single_dst_page) { + int const first_page_token = first_dst_page * params.page_size; + int const logical_page = batch * params.max_num_pages_per_seq + first_dst_page; + int const phys_page = params.ptr_page_table[logical_page]; + single_row_base = phys_page * params.page_size + (cache_len_old - first_page_token); + } + } + + if constexpr (PagedKV) { + if (!single_dst_page) { + // Cross-page append is common for large k_new; resolve the page + // table once per destination page instead of once per token. + for (int new_tok0 = 0; new_tok0 < new_len;) { + int const dst_tok0 = cache_len_old + new_tok0; + int const page = dst_tok0 / params.page_size; + int const tok_in_page0 = dst_tok0 - page * params.page_size; + int const tokens_in_page = params.page_size - tok_in_page0; + int const page_len = tokens_in_page < (new_len - new_tok0) ? tokens_in_page : (new_len - new_tok0); + int const logical_page = batch * params.max_num_pages_per_seq + page; + int const phys_page = params.ptr_page_table[logical_page]; + int const row_base = phys_page * params.page_size + tok_in_page0; + size_t src_base = + ((size_t)(new_begin + new_tok0) * (size_t)num_heads_kv + (size_t)kv_head) * (size_t)head_size; + size_t const token_stride = (size_t)num_heads_kv * (size_t)head_size; + + if (flatten_token_vectors) { + for (int page_vec = worker; page_vec < page_len * vecs_per_token; page_vec += worker_count) { + int const page_tok = page_vec / vecs_per_token; + int const d_vec = page_vec - page_tok * vecs_per_token; + int const dst_row = row_base + page_tok; + size_t const token_src_base = src_base + (size_t)page_tok * token_stride; + int const d = d_vec * kVecElems; + size_t const src = token_src_base + (size_t)d; + StoreVec const k_value = *reinterpret_cast(params.ptr_K_new + src); + StoreVec const v_value = *reinterpret_cast(params.ptr_V_new + src); + *reinterpret_cast(&K_dst(dst_row, d)) = k_value; + *reinterpret_cast(&V_dst(d, dst_row)) = v_value; + } + } else { + for (int page_tok = sub_group_id; page_tok < page_len; page_tok += active_sg_count) { + int const dst_row = row_base + page_tok; + size_t const token_src_base = src_base + (size_t)page_tok * token_stride; + for (int d_vec = lane_idx; d_vec < vecs_per_token; d_vec += intel::sg_size) { + int const d = d_vec * kVecElems; + size_t const src = token_src_base + (size_t)d; + StoreVec const k_value = *reinterpret_cast(params.ptr_K_new + src); + StoreVec const v_value = *reinterpret_cast(params.ptr_V_new + src); + *reinterpret_cast(&K_dst(dst_row, d)) = k_value; + *reinterpret_cast(&V_dst(d, dst_row)) = v_value; + } + } + } + new_tok0 += page_len; + } + return true; + } + } + + if (flatten_token_vectors) { + for (int new_vec = worker; new_vec < new_len * vecs_per_token; new_vec += worker_count) { + int const new_tok = new_vec / vecs_per_token; + int const d_vec = new_vec - new_tok * vecs_per_token; + int const new_abs_tok = new_begin + new_tok; + int const dst_row = single_row_base + new_tok; + size_t const src_base = ((size_t)new_abs_tok * (size_t)num_heads_kv + (size_t)kv_head) * (size_t)head_size; + int const d = d_vec * kVecElems; + size_t const src = src_base + (size_t)d; + StoreVec const k_value = *reinterpret_cast(params.ptr_K_new + src); + StoreVec const v_value = *reinterpret_cast(params.ptr_V_new + src); + *reinterpret_cast(&K_dst(dst_row, d)) = k_value; + *reinterpret_cast(&V_dst(d, dst_row)) = v_value; + } + } else { + for (int new_tok = sub_group_id; new_tok < new_len; new_tok += active_sg_count) { + int const new_abs_tok = new_begin + new_tok; + int const dst_row = single_row_base + new_tok; + size_t const src_base = ((size_t)new_abs_tok * (size_t)num_heads_kv + (size_t)kv_head) * (size_t)head_size; + for (int d_vec = lane_idx; d_vec < vecs_per_token; d_vec += intel::sg_size) { + int const d = d_vec * kVecElems; + size_t const src = src_base + (size_t)d; + StoreVec const k_value = *reinterpret_cast(params.ptr_K_new + src); + StoreVec const v_value = *reinterpret_cast(params.ptr_V_new + src); + *reinterpret_cast(&K_dst(dst_row, d)) = k_value; + *reinterpret_cast(&V_dst(d, dst_row)) = v_value; + } + } + } + return true; + } + + CUTLASS_DEVICE + void store_kv_new( + TensorK_cache2D const& K_cache_2D, + TensorV_cache2D const& V_cache_2D, + int batch, + int kv_head, + int num_heads_kv, + int thr_id, + int append_store_len = -1, + int append_store_begin = 0) const { + if constexpr (AppendKV) { + int const new_len_total = get_kv_new_len(batch); + int const store_begin = cute::min(cute::max(append_store_begin, 0), new_len_total); + int const available_len = new_len_total - store_begin; + int const new_len = + append_store_len < 0 ? available_len : (append_store_len < available_len ? append_store_len : available_len); + if (new_len <= 0) { + return; + } + + auto& K_dst = const_cast(K_cache_2D); + auto& V_dst = const_cast(V_cache_2D); + int const lane_idx = thr_id % intel::sg_size; + int const sub_group_id = thr_id / intel::sg_size; + int const new_begin = (params.ptr_cu_seqlens_k_new != nullptr ? params.ptr_cu_seqlens_k_new[batch] + : batch * params.seq_len_kv_new) + + store_begin; + int const cache_len_old = params.ptr_kv_cache_seqlens[batch] + store_begin; + int const head_size_qk = size<1>(K_cache_2D); + int const head_size_vo = size<0>(V_cache_2D); + int const max_hd = head_size_qk > head_size_vo ? head_size_qk : head_size_vo; + + // Every WG that can read this appended KV range scatters the same source + // k_new/v_new values before loading K/V. Multiple WGs, including GQA + // query heads that map to one kv_head, therefore overwrite identical bit + // patterns at the same cache address; the append is intentionally + // idempotent and does not rely on a grid-wide producer. + if constexpr ( + sizeof(typename TensorK_cache::element_type) == 2 && sizeof(typename TensorV_cache::element_type) == 2) { + if constexpr (AppendWideStore) { + if (store_kv_new_vectorized<16, cutlass::ulonglong4>( + K_cache_2D, + V_cache_2D, + batch, + kv_head, + num_heads_kv, + lane_idx, + sub_group_id, + new_len, + new_begin, + cache_len_old, + head_size_qk, + head_size_vo)) { + return; + } + } else { + if (store_kv_new_vectorized<8, cutlass::ulonglong2>( + K_cache_2D, + V_cache_2D, + batch, + kv_head, + num_heads_kv, + lane_idx, + sub_group_id, + new_len, + new_begin, + cache_len_old, + head_size_qk, + head_size_vo)) { + return; + } + } + } + + for (int linear = lane_idx; linear < new_len * max_hd; linear += intel::sg_size) { + int const d = linear % max_hd; + int const new_tok = linear / max_hd; + int const new_abs_tok = new_begin + new_tok; + int const dst_tok = cache_len_old + new_tok; + int dst_row = dst_tok; + if constexpr (PagedKV) { + int const page = dst_tok / params.page_size; + int const tok_in_page = dst_tok - page * params.page_size; + int const logical_page = batch * params.max_num_pages_per_seq + page; + int const phys_page = params.ptr_page_table[logical_page]; + dst_row = phys_page * params.page_size + tok_in_page; + } + + if (d < head_size_qk) { + size_t const src = + ((size_t)new_abs_tok * (size_t)num_heads_kv + (size_t)kv_head) * (size_t)head_size_qk + (size_t)d; + K_dst(dst_row, d) = params.ptr_K_new[src]; + } + if (d < head_size_vo) { + size_t const src = + ((size_t)new_abs_tok * (size_t)num_heads_kv + (size_t)kv_head) * (size_t)head_size_vo + (size_t)d; + V_dst(d, dst_row) = params.ptr_V_new[src]; + } + } + } else { + (void)K_cache_2D; + (void)V_cache_2D; + (void)batch; + (void)kv_head; + (void)num_heads_kv; + (void)thr_id; + (void)append_store_len; + (void)append_store_begin; + } + } + template CUTLASS_DEVICE void operator()( TensorQ2D const& Q_2D, // (q,d) @@ -271,8 +569,12 @@ struct FMHAFwdMainloop< int seq_len, int seq_len_kv_cache, int l_coord, + int kv_head, + int num_heads_kv, int full_tile_offset, int discard_seq_coord, + int append_store_len = -1, + int append_store_begin = 0, TensorK_cache2D const& K_cache_2D = TensorK_cache2D{}, TensorV_cache2D const& V_cache_2D = TensorV_cache2D{}, float scale_k = 1.0f) { // FP8 K per-tensor dequant scale @@ -365,20 +667,25 @@ struct FMHAFwdMainloop< // Kernel // ------ - /* Initialization steps for first block: Q/K prefetch, O init */ - /* TODO: limit D prefetch for large head size, and reorder K prefetches */ - int kblocks_cache = ceil_div(seq_len_kv_cache, get<1>(TileShapeQK{})); - int page_idx = blk_k0; - int next_page_idx = blk_k0; - if constexpr (PagedKV) { - next_page_idx = get_physical_k_tile(blk_k0, l_coord, seq_len_kv_cache); + if constexpr (AppendKV) { + store_kv_new( + K_cache_2D, V_cache_2D, l_coord, kv_head, num_heads_kv, thr_id, append_store_len, append_store_begin); + bool direct_append = false; + if constexpr (AppendDirectLoad) { + constexpr int kTileKV = get<1>(TileShapeQK{}); + int const cache_len_old = params.ptr_kv_cache_seqlens[l_coord]; + int const new_begin = params.ptr_cu_seqlens_k_new != nullptr ? params.ptr_cu_seqlens_k_new[l_coord] + : l_coord * params.seq_len_kv_new; + direct_append = (cache_len_old % kTileKV) == 0 && (new_begin % kTileKV) == 0; + } + if (!direct_append) { + barrier(); + } } + for (int D = 0; D < size<3>(pQgQ); D++) { prefetch(prefetch_q, pQgQ(_, _, _, D)); } - for (int D = 0; D < size<4>(pKgK); D++) { - prefetch(prefetch_k_cache, pKgK_cache(_, _, _, next_page_idx, D)); - } // Always initialize the per-WG accumulators: the caller (kernel) may pass // blk_k0 > 0 when sliding-window pruning skips leading K blocks, so we can // no longer key initialization off of (blk_k0 == 0). @@ -398,126 +705,185 @@ struct FMHAFwdMainloop< qk_scale = params.scale * static_cast(scale_k); } - /* Main loop, blocked in k. */ - for (int K = blk_k0; K < blk_k1 && K < kblocks_cache; K++) { - /* Split barrier to keep threads together */ - barrier_arrive(ScopeWorkgroup); + constexpr int kTileKV = get<1>(TileShapeQK{}); + int const kblocks_total = ceil_div(seq_len_kv_cache, kTileKV); + bool direct_append = false; + int direct_block0 = kblocks_total; + int direct_source_block0 = 0; + if constexpr (AppendDirectLoad) { + int const cache_len_old = params.ptr_kv_cache_seqlens[l_coord]; + int const new_begin = params.ptr_cu_seqlens_k_new != nullptr ? params.ptr_cu_seqlens_k_new[l_coord] + : l_coord * params.seq_len_kv_new; + direct_append = (cache_len_old % kTileKV) == 0 && (new_begin % kTileKV) == 0; + direct_block0 = cache_len_old / kTileKV; + direct_source_block0 = new_begin / kTileKV; + } - bool need_causal = false; - if constexpr (CausalMask) { - need_causal = K >= blk_k1_causal; + auto run_k_blocks = [&](auto source_tag, int loop_k0, int loop_k1) { + constexpr bool LoadDirect = decltype(source_tag)::value; + if (loop_k0 >= loop_k1) { + return; } - page_idx = next_page_idx; - next_page_idx = K + 1; - if constexpr (PagedKV) { - next_page_idx = get_physical_k_tile(next_page_idx, l_coord, seq_len_kv_cache); + int next_source_idx = loop_k0; + if constexpr (LoadDirect) { + next_source_idx = direct_source_block0 + loop_k0 - direct_block0; + for (int D = 0; D < size<4>(pKgK); D++) { + prefetch(prefetch_k, pKgK(_, _, _, next_source_idx, D)); + } + } else { + if constexpr (PagedKV) { + next_source_idx = get_physical_k_tile(loop_k0, l_coord, seq_len_kv_cache); + } + for (int D = 0; D < size<4>(pKgK); D++) { + prefetch(prefetch_k_cache, pKgK_cache(_, _, _, next_source_idx, D)); + } } - /* GEMM 1: S = K * Q */ - clear(tSrS); - CUTLASS_PRAGMA_UNROLL - for (int D = 0; D < size<4>(tKgK); D++) { - copy(copy_q, tQgQ(_, _, _, D), tQrQ); - copy(copy_k_cache, tKgK_cache(_, _, _, page_idx, D), tKrK); - reorder(tQrQ, tSrQ); - reorder(tKrK, tSrK); - cute::gemm(mma_qk, tSrQ, tSrK, tSrS); - } + for (int K = loop_k0; K < loop_k1 && K < kblocks_total; K++) { + barrier_arrive(ScopeWorkgroup); - /* V prefetch for GEMM 2 */ - CUTLASS_PRAGMA_UNROLL - for (int VV = 0; VV < VTiles; VV++) { - prefetch(prefetch_v_cache, pVgV_cache(_, _, _, VV, page_idx)); - } + bool need_causal = false; + if constexpr (CausalMask) { + need_causal = K >= blk_k1_causal; + } - /* Causal masking */ - if constexpr (CausalMask) { - if (need_causal) { - /* Masking scalars */ - // TODO: use a more general code path for causal masking. - int lane_id = thr_id % intel::sg_size; - constexpr int sg_tile_q = get<0>(TileShapeQK{}) / SGPerWG::value; - int row_base = get<0>(blk_qv) * get<0>(TileShapeQK{}) + (thr_id / intel::sg_size) * sg_tile_q; - - constexpr int kTileK = get<1>(TileShapeQK{}); - constexpr int n_reps = kTileK / intel::sg_size; - constexpr int elems_per_n = tSrS.size() / n_reps; - int k_base = K * kTileK; - CUTLASS_PRAGMA_UNROLL - for (int n = 0; n < n_reps; n++) { - int col = k_base + n * intel::sg_size + lane_id; - int causal_bound = col - full_tile_offset - row_base; - CUTLASS_PRAGMA_UNROLL - for (int j = 0; j < elems_per_n; j++) { - if (j < causal_bound) { - tSrS(n * elems_per_n + j) = ElementS(-INFINITY); - } + int const source_idx = next_source_idx; + int const next_k = K + 1; + bool const has_next_k = next_k < loop_k1 && next_k < kblocks_total; + if constexpr (LoadDirect) { + next_source_idx = direct_source_block0 + next_k - direct_block0; + } else { + next_source_idx = next_k; + if constexpr (PagedKV) { + if (has_next_k) { + next_source_idx = get_physical_k_tile(next_k, l_coord, seq_len_kv_cache); } } } - } - /* Local/sliding window masking */ - if constexpr (LocalMask) { - Tensor cPgP = make_identity_tensor(make_shape(seq_len, seq_len)); - Tensor gP = local_tile(cPgP, take<0, 2>(TileShapeQK{}), make_coord(get<0>(blk_qv), K)); - auto cS_thread = thr_mma_qk.partition_C(gP); + clear(tSrS); CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < tSrS.size(); ++i) { - int row_idx = get<0>(cS_thread(i)); - int col_idx = get<1>(cS_thread(i)); - // PackGQA decode: every packed M row is the same decode token, so the - // KV position is full_tile_offset regardless of the per-row (head) - // index. Non-packed keeps the per-row sequence position. - int row_kv_idx = (PackGQA_ ? 0 : row_idx) + full_tile_offset; - bool left_mask = col_idx < row_kv_idx - params.window_size_left; - bool right_mask = col_idx > row_kv_idx + params.window_size_right; - if (left_mask || right_mask) { - tSrS(i) = ElementS(-INFINITY); + for (int D = 0; D < size<4>(tKgK); D++) { + copy(copy_q, tQgQ(_, _, _, D), tQrQ); + if constexpr (LoadDirect) { + copy(copy_k, tKgK(_, _, _, source_idx, D), tKrK); + } else { + copy(copy_k_cache, tKgK_cache(_, _, _, source_idx, D), tKrK); } + reorder(tQrQ, tSrQ); + reorder(tKrK, tSrK); + cute::gemm(mma_qk, tSrQ, tSrK, tSrS); } - } - /* k masking for remainder tiles */ - if (check_remainder_k && K == total_blk - 1) { - FragSCol k_rem_mask; - int k_val = get<0>(tKgK_cache(0, 0, 0, K, 0)); - int k = k_val + get_sub_group().get_local_id()[0]; CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < k_rem_mask.size(); i++, k += intel::sg_size) { - k_rem_mask(i) = (k < seq_len) ? ElementS(sycl::nan(0u)) : ElementS(-INFINITY); + for (int VV = 0; VV < VTiles; VV++) { + if constexpr (LoadDirect) { + prefetch(prefetch_v, pVgV(_, _, _, VV, source_idx)); + } else { + prefetch(prefetch_v_cache, pVgV_cache(_, _, _, VV, source_idx)); + } } - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < tSrS.size(); i++) { - tSrS(i) = sycl::fmin(tSrS(i), broadcast<1>(k_rem_mask, tSrS, i)); + + if constexpr (CausalMask) { + if (need_causal) { + /* Masking scalars */ + // TODO: use a more general code path for causal masking. + int lane_id = thr_id % intel::sg_size; + constexpr int sg_tile_q = get<0>(TileShapeQK{}) / SGPerWG::value; + int row_base = get<0>(blk_qv) * get<0>(TileShapeQK{}) + (thr_id / intel::sg_size) * sg_tile_q; + + constexpr int kTileK = get<1>(TileShapeQK{}); + constexpr int n_reps = kTileK / intel::sg_size; + int const elems_per_n = tSrS.size() / n_reps; + int k_base = K * kTileK; + CUTLASS_PRAGMA_UNROLL + for (int n = 0; n < n_reps; n++) { + int col = k_base + n * intel::sg_size + lane_id; + int causal_bound = col - full_tile_offset - row_base; + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < elems_per_n; j++) { + if (j < causal_bound) { + tSrS(n * elems_per_n + j) = ElementS(-INFINITY); + } + } + } + } } - } - /* Apply softmax and scaling (tA rescaling fused into GEMM2 VTile loop) */ - auto rescale = softmax(K == blk_k0, tSrS, tA_max, tA_sum, qk_scale); - reorder(tSrS, tArP); + if constexpr (LocalMask) { + Tensor cPgP = make_identity_tensor(make_shape(seq_len, seq_len)); + Tensor gP = local_tile(cPgP, take<0, 2>(TileShapeQK{}), make_coord(get<0>(blk_qv), K)); + auto cS_thread = thr_mma_qk.partition_C(gP); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < tSrS.size(); ++i) { + int row_idx = get<0>(cS_thread(i)); + int col_idx = get<1>(cS_thread(i)); + int row_kv_idx = (PackGQA_ ? 0 : row_idx) + full_tile_offset; + bool left_mask = col_idx < row_kv_idx - params.window_size_left; + bool right_mask = col_idx > row_kv_idx + params.window_size_right; + if (left_mask || right_mask) { + tSrS(i) = ElementS(-INFINITY); + } + } + } - /* GEMM 2: A += P * V, split in v dimension. */ - CUTLASS_PRAGMA_UNROLL - for (int VV = 0; VV < VTiles; VV++) { - copy(copy_v_cache, tVgV_cache(_, _, _, VV, page_idx), tVrV); - reorder(tVrV, tArV); - if (K != blk_k0) { + if (check_remainder_k && K == total_blk - 1) { + FragSCol k_rem_mask; + int k_val = get<0>(tKgK_cache(0, 0, 0, K, 0)); + int k = k_val + get_sub_group().get_local_id()[0]; + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < k_rem_mask.size(); i++, k += intel::sg_size) { + k_rem_mask(i) = (k < seq_len) ? ElementS(sycl::nan(0u)) : ElementS(-INFINITY); + } CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < tArA.size() / VTiles; i++) { - tArA(_, _, _, VV)(i) *= broadcast<0>(rescale, tArA, i); + for (int i = 0; i < tSrS.size(); i++) { + tSrS(i) = sycl::fmin(tSrS(i), broadcast<1>(k_rem_mask, tSrS, i)); } } - cute::gemm(mma_pv, tArP, tArV, tArA(_, _, _, VV)); - } - /* K prefetch */ - for (int D = 0; D < size<4>(pKgK); D++) { - prefetch(prefetch_k_cache, pKgK_cache(_, _, _, next_page_idx, D)); + auto rescale = softmax(K == blk_k0, tSrS, tA_max, tA_sum, qk_scale); + reorder(tSrS, tArP); + + CUTLASS_PRAGMA_UNROLL + for (int VV = 0; VV < VTiles; VV++) { + if constexpr (LoadDirect) { + copy(copy_v, tVgV(_, _, _, VV, source_idx), tVrV); + } else { + copy(copy_v_cache, tVgV_cache(_, _, _, VV, source_idx), tVrV); + } + reorder(tVrV, tArV); + if (K != blk_k0) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < tArA.size() / VTiles; i++) { + tArA(_, _, _, VV)(i) *= broadcast<0>(rescale, tArA, i); + } + } + cute::gemm(mma_pv, tArP, tArV, tArA(_, _, _, VV)); + } + + if (has_next_k) { + if constexpr (LoadDirect) { + for (int D = 0; D < size<4>(pKgK); D++) { + prefetch(prefetch_k, pKgK(_, _, _, next_source_idx, D)); + } + } else { + for (int D = 0; D < size<4>(pKgK); D++) { + prefetch(prefetch_k_cache, pKgK_cache(_, _, _, next_source_idx, D)); + } + } + } + + barrier_wait(ScopeWorkgroup); } + }; - barrier_wait(ScopeWorkgroup); + int const cache_loop_k1 = direct_append ? cute::min(blk_k1, direct_block0) : blk_k1; + run_k_blocks(cute::false_type{}, blk_k0, cache_loop_k1); + if constexpr (AppendDirectLoad) { + if (direct_append) { + run_k_blocks(cute::true_type{}, cute::max(blk_k0, direct_block0), blk_k1); + } } } diff --git a/src/sycl/kernels/flash_attention_v2/kernel/xe_fmha_append_kv.hpp b/src/sycl/kernels/flash_attention_v2/kernel/xe_fmha_append_kv.hpp new file mode 100644 index 000000000..fb0ad5fce --- /dev/null +++ b/src/sycl/kernels/flash_attention_v2/kernel/xe_fmha_append_kv.hpp @@ -0,0 +1,156 @@ +/*************************************************************************************************** + * Copyright (C) 2026 Intel Corporation, All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + **************************************************************************************************/ +#pragma once + +#include "cute/tensor.hpp" +#include "cute/util/type_traits.hpp" +#include "cutlass/cutlass.h" +#include "cutlass/kernel_hardware_info.hpp" +#include "cutlass/numeric_types.h" + +namespace cutlass::fmha::kernel { + +// Scatter a large appended KV range once before attention. The fused mainloop +// path remains preferable for short appends because it avoids a second launch. +template +class XeFMHAAppendKVKernel { + public: + using ElementK = ElementK_; + using ElementV = ElementV_; + static constexpr int NumSGs = NumSGs_; + static constexpr int SgSize = cute::intel::sg_size; + static constexpr int ThreadsPerWG = NumSGs * SgSize; + static constexpr int TokensPerWG = NumSGs; + + struct Arguments { + ElementK* ptr_K_cache = nullptr; + ElementV* ptr_V_cache = nullptr; + ElementK const* ptr_K_new = nullptr; + ElementV const* ptr_V_new = nullptr; + int const* ptr_cu_seqlens_k_new = nullptr; + int seq_len_kv_new = 0; + int const* ptr_cache_seqlens = nullptr; + int const* ptr_page_table = nullptr; + int page_size = 0; + int max_num_pages_per_seq = 0; + int batch = 0; + int num_heads_kv = 0; + int head_size_qk = 0; + int head_size_vo = 0; + int max_seq_len_kv_new = 0; + bool const* skip_batch_mask = nullptr; + }; + + using Params = Arguments; + struct SharedStorage {}; + static constexpr int SharedStorageSize = 0; + + static Params to_underlying_arguments(Arguments const& args, void*) { + return args; + } + + static bool can_implement(Arguments const& args) { + return args.ptr_K_cache != nullptr && args.ptr_V_cache != nullptr && + args.ptr_K_new != nullptr && args.ptr_V_new != nullptr && + args.ptr_cache_seqlens != nullptr && args.ptr_page_table != nullptr && + args.batch > 0 && args.num_heads_kv > 0 && args.head_size_qk > 0 && + args.head_size_vo > 0 && args.max_seq_len_kv_new > 0 && + (args.ptr_cu_seqlens_k_new != nullptr || args.seq_len_kv_new > 0) && + args.page_size > 0 && args.max_num_pages_per_seq > 0; + } + + static int get_workspace_size(Arguments const&) { + return 0; + } + + static cutlass::Status initialize_workspace( + Arguments const&, void* = nullptr, void* = nullptr, void* = nullptr) { + return Status::kSuccess; + } + + static dim3 get_grid_shape(Params const& params) { + return dim3( + uint32_t(cute::ceil_div(params.max_seq_len_kv_new, TokensPerWG)), + uint32_t(params.num_heads_kv), + uint32_t(params.batch)); + } + + static dim3 get_block_shape() { + return dim3(ThreadsPerWG, 1, 1); + } + + CUTLASS_DEVICE + void operator()(Params const& params, char*) const { + int const batch_idx = int(BlockIdxZ()); + if (params.skip_batch_mask != nullptr && params.skip_batch_mask[batch_idx]) return; + + int new_begin; + int new_len; + if (params.ptr_cu_seqlens_k_new != nullptr) { + new_begin = params.ptr_cu_seqlens_k_new[batch_idx]; + new_len = params.ptr_cu_seqlens_k_new[batch_idx + 1] - new_begin; + } else { + new_len = params.seq_len_kv_new; + new_begin = batch_idx * new_len; + } + + int const token = int(BlockIdxX()) * TokensPerWG + int(ThreadIdxX()) / SgSize; + if (token >= new_len) return; + + int const lane = int(ThreadIdxX()) % SgSize; + int const kv_head = int(BlockIdxY()); + int const dst_token = params.ptr_cache_seqlens[batch_idx] + token; + int const page = dst_token / params.page_size; + int const page_token = dst_token - page * params.page_size; + int const dst_row = + params.ptr_page_table[batch_idx * params.max_num_pages_per_seq + page] * + params.page_size + + page_token; + int const src_token = new_begin + token; + + size_t const k_src = + (size_t(src_token) * params.num_heads_kv + kv_head) * params.head_size_qk; + size_t const v_src = + (size_t(src_token) * params.num_heads_kv + kv_head) * params.head_size_vo; + size_t const k_dst = + (size_t(dst_row) * params.num_heads_kv + kv_head) * params.head_size_qk; + size_t const v_dst = + (size_t(dst_row) * params.num_heads_kv + kv_head) * params.head_size_vo; + + copy_row(params.ptr_K_cache, params.ptr_K_new, k_dst, k_src, params.head_size_qk, lane); + copy_row(params.ptr_V_cache, params.ptr_V_new, v_dst, v_src, params.head_size_vo, lane); + } + + private: + template + CUTLASS_DEVICE static void copy_row( + Element* dst, Element const* src, size_t dst_offset, size_t src_offset, int elements, int lane) { + using Vec32 = cute::intel::uint8; + using Vec16 = cutlass::ulonglong2; + if constexpr (sizeof(Element) == 2) { + if ((elements % 16) == 0) { + for (int i = lane; i < elements / 16; i += SgSize) { + int const offset = i * 16; + *reinterpret_cast(dst + dst_offset + offset) = + *reinterpret_cast(src + src_offset + offset); + } + return; + } + if ((elements % 8) == 0) { + for (int i = lane; i < elements / 8; i += SgSize) { + int const offset = i * 8; + *reinterpret_cast(dst + dst_offset + offset) = + *reinterpret_cast(src + src_offset + offset); + } + return; + } + } + for (int d = lane; d < elements; d += SgSize) { + dst[dst_offset + d] = src[src_offset + d]; + } + } +}; + +} // namespace cutlass::fmha::kernel diff --git a/src/sycl/kernels/flash_attention_v2/kernel/xe_fmha_fwd_kernel.hpp b/src/sycl/kernels/flash_attention_v2/kernel/xe_fmha_fwd_kernel.hpp index 1053a4559..54bc98036 100644 --- a/src/sycl/kernels/flash_attention_v2/kernel/xe_fmha_fwd_kernel.hpp +++ b/src/sycl/kernels/flash_attention_v2/kernel/xe_fmha_fwd_kernel.hpp @@ -275,6 +275,13 @@ class XeFMHAFwdKernel { auto sequence_length_shape = get_sequence_length_shape(s, idx_b); auto [seq_len_qo, seq_len_kv, seq_len_kv_cache] = sequence_length_shape; + int seq_k_eff = seq_len_kv_cache; + if constexpr (CollectiveMainloop::AppendKV) { + int const seq_k_new = CollectiveMainloop::get_kv_new_len(params.mainloop, idx_b); + if (seq_k_new > 0) { + seq_k_eff = params.mainloop.ptr_kv_cache_seqlens[idx_b] + seq_k_new; + } + } // M extent of the Q/O tile: the packed GQA group for decode, otherwise the // query sequence length. Masking below still uses the real seq_len_qo so // the decode KV position (seq_len_kv_cache - seq_len_qo) stays correct. @@ -286,7 +293,7 @@ class XeFMHAFwdKernel { // auto offset = cute::min(seq_len_qo, seq_len_kv_cache); auto offset = seq_len_qo; auto discard_seq_coord = seq_len_qo - offset; - auto full_tile_offset = seq_len_kv_cache - offset; + auto full_tile_offset = seq_k_eff - offset; int seq_coord = cute::min(seq_len_qo, (blk_q * get<0>(TileShapeQK{}) + q_offset_sg)); // if (CollectiveMainloop::CausalMask && seq_coord < discard_seq_coord) continue; @@ -297,12 +304,47 @@ class XeFMHAFwdKernel { // const int seq_len = seq_len_new + seq_len_kv_cache; // const int k_blocks = cute::ceil_div(seq_len, get<1>(TileShapeQK{})); - const int seq_len = CollectiveMainloop::CausalMask - ? cute::min(seq_len_kv_cache, full_tile_offset + seq_coord + q_sg_tile) - : seq_len_kv_cache; + const int seq_len = + CollectiveMainloop::CausalMask ? cute::min(seq_k_eff, full_tile_offset + seq_coord + q_sg_tile) : seq_k_eff; const int k_blocks = cute::ceil_div(seq_len, get<1>(TileShapeQK{})); const int k_blocks_causal = CollectiveMainloop::CausalMask ? (seq_coord + full_tile_offset) / get<1>(TileShapeQK{}) : 0; + int append_store_len = -1; + int append_store_begin = 0; + if constexpr (CollectiveMainloop::AppendKV) { + int const seq_k_new = CollectiveMainloop::get_kv_new_len(params.mainloop, idx_b); + if (seq_k_new > 0) { + append_store_len = seq_k_new; + if constexpr (CollectiveMainloop::AppendDirectLoad) { + int const cache_len_old = params.mainloop.ptr_kv_cache_seqlens[idx_b]; + int const new_begin = params.mainloop.ptr_cu_seqlens_k_new != nullptr + ? params.mainloop.ptr_cu_seqlens_k_new[idx_b] + : idx_b * params.mainloop.seq_len_kv_new; + constexpr int kTileKV = get<1>(TileShapeQK{}); + bool const direct_batch = (cache_len_old % kTileKV) == 0 && (new_begin % kTileKV) == 0; + if (direct_batch) { + if (get<0>(blk_qv) == 0 && get<1>(blk_qv) == 0) { + int const writer = head_q % head_group_q; + append_store_begin = seq_k_new * writer / head_group_q; + int const append_store_end = seq_k_new * (writer + 1) / head_group_q; + append_store_len = append_store_end - append_store_begin; + } else { + append_store_len = 0; + } + } + } + if constexpr (CollectiveMainloop::CausalMask && !PackGQA_) { + // Without a grid-wide barrier, each WG must write every appended + // token it may read; causal tiles only need the visible prefix. + int const cache_len_old = params.mainloop.ptr_kv_cache_seqlens[idx_b]; + int const tile_q = get<0>(TileShapeQK{}); + int const q_tile_end = cute::min(seq_len_qo, (blk_q + 1) * tile_q); + int const visible_k_end = cute::min(seq_k_eff, full_tile_offset + q_tile_end); + append_store_len = cute::max(0, visible_k_end - cache_len_old); + append_store_len = cute::min(append_store_len, seq_k_new); + } + } + } // Sliding-window pruning: skip K blocks that are entirely outside the // [row - window_size_left, row + window_size_right] band for all rows in @@ -353,15 +395,25 @@ class XeFMHAFwdKernel { // contiguous region, so the extent must be this batch's KV length to keep the // 2D block loads in-bounds. int kv_seq_extent = CollectiveMainloop::PagedKV ? int(s.seq_len_kv_cache.total_length) : int(seq_len_kv_cache); + int kv_input_extent = kv_seq_extent; + bool use_direct_append_source = false; + if constexpr (CollectiveMainloop::AppendDirectLoad) { + use_direct_append_source = true; + kv_input_extent = int(s.seq_len_kv.total_length); + } // PackGQA folds the head_group_q query heads into M and grids over KV // heads, so the Q/O head extent collapses to num_heads_kv. auto q_head_count = PackGQA_ ? s.num_heads_kv : s.num_heads_q; auto shape_Q = make_shape(m_extent, s.head_size_qk, q_head_count, batch_dim); - auto shape_K = make_shape(kv_seq_extent, s.head_size_qk, s.num_heads_kv, batch_dim); - auto shape_V = make_shape(s.head_size_vo, kv_seq_extent, s.num_heads_kv, batch_dim); + auto shape_K = make_shape(kv_input_extent, s.head_size_qk, s.num_heads_kv, batch_dim); + auto shape_V = make_shape(s.head_size_vo, kv_input_extent, s.num_heads_kv, batch_dim); + auto shape_K_cache = make_shape(kv_seq_extent, s.head_size_qk, s.num_heads_kv, batch_dim); + auto shape_V_cache = make_shape(s.head_size_vo, kv_seq_extent, s.num_heads_kv, batch_dim); auto shape_O = make_shape(m_extent, s.head_size_vo, q_head_count, batch_dim); auto dcQ = const_cast(p.Q + offset_q); + auto dcK = const_cast(use_direct_append_source ? p.K : p.K_cache); + auto dcV = const_cast(use_direct_append_source ? p.V : p.V_cache); auto dcK_cache = const_cast(p.K_cache + offset_k_cache); auto dcV_cache = const_cast(p.V_cache + offset_v_cache); auto dcO = const_cast(p.O + offset_o); @@ -369,13 +421,19 @@ class XeFMHAFwdKernel { auto layout_q = is_var_len ? make_ordered_layout(shape_Q, VarLenQLayoutStep_{}) : make_layout(shape_Q, p.dQ); auto layout_k = is_var_len ? make_ordered_layout(shape_K, VarLenKLayoutStep_{}) : make_layout(shape_K, p.dK); auto layout_v = is_var_len ? make_ordered_layout(shape_V, VarLenVLayoutStep_{}) : make_layout(shape_V, p.dV); + auto layout_k_cache = is_var_len ? make_ordered_layout(shape_K_cache, VarLenKLayoutStep_{}) + : make_layout(shape_K_cache, p.dK_cache); + auto layout_v_cache = is_var_len ? make_ordered_layout(shape_V_cache, VarLenVLayoutStep_{}) + : make_layout(shape_V_cache, p.dV_cache); // NHD layout for GQA auto layout_o = is_var_len ? make_ordered_layout(shape_O, VarLenOLayoutStep_{}) : make_layout(shape_O, p.dO); Tensor Q = make_tensor(make_gmem_ptr(dcQ), layout_q); - Tensor K_cache = make_tensor(make_gmem_ptr(dcK_cache), layout_k); - Tensor V_cache = make_tensor(make_gmem_ptr(dcV_cache), layout_v); + Tensor K = make_tensor(make_gmem_ptr(dcK), layout_k); + Tensor V = make_tensor(make_gmem_ptr(dcV), layout_v); + Tensor K_cache = make_tensor(make_gmem_ptr(dcK_cache), layout_k_cache); + Tensor V_cache = make_tensor(make_gmem_ptr(dcV_cache), layout_v_cache); Tensor O = make_tensor(make_gmem_ptr(dcO), layout_o); // O accumulator types FragA tArA; @@ -398,8 +456,8 @@ class XeFMHAFwdKernel { CollectiveMainloop mainloop(params.mainloop, shared_storage.mainloop); mainloop( Q(_, _, q_head_idx, l_coord), - K_cache(_, _, head, l_coord), - V_cache(_, _, head, l_coord), + K(_, _, head, l_coord), + V(_, _, head, l_coord), tArA, tA_max, tA_sum, @@ -410,10 +468,14 @@ class XeFMHAFwdKernel { k_blocks_causal, thr_id, seq_len, - seq_len_kv_cache, + seq_k_eff, idx_b, + head, + s.num_heads_kv, full_tile_offset, discard_seq_coord, + append_store_len, + append_store_begin, K_cache(_, _, head, l_coord), V_cache(_, _, head, l_coord), scale_k); diff --git a/src/sycl/kernels/flash_attention_v2/xe_fmha_fwd_prefill_kernel.cpp.in b/src/sycl/kernels/flash_attention_v2/xe_fmha_fwd_prefill_kernel.cpp.in index 57f33c0a0..10c633d1e 100644 --- a/src/sycl/kernels/flash_attention_v2/xe_fmha_fwd_prefill_kernel.cpp.in +++ b/src/sycl/kernels/flash_attention_v2/xe_fmha_fwd_prefill_kernel.cpp.in @@ -36,6 +36,20 @@ namespace prefill { +namespace { + +template +void run_paged_tile(const Arguments& params) { + AT_DISPATCH_BOOL_NO_RETURN(params.is_local, LocalMask, { + AT_DISPATCH_BOOL_NO_RETURN(params.is_causal, Causal, { + FMHAConfig::run_paged( + params); + }); + }); +} + +} // namespace + template <> void FmhaPrefillRunner<@HEAD_DIM@>::operator()(const Arguments& params) const { bool use_sink = (params.softmax_sink_ptr != nullptr); @@ -49,6 +63,193 @@ void FmhaPrefillRunner<@HEAD_DIM@>::operator()(const Arguments& params) const { using SubgroupLayoutQK = cute::Layout, cute::_1, cute::_1>, cute::Stride>; +#if @APPENDKV_Q256_TILE@ + if (params.total_kvnew > 0 && params.seqlen_q == @APPENDKV_Q256_Q@ && + params.seqlen_kvnew == @APPENDKV_Q256_Q@ && + params.total_q == params.b * @APPENDKV_Q256_Q@ && + params.total_kvnew == params.b * @APPENDKV_Q256_Q@ && !params.is_causal && !params.is_local && !use_sink && + params.d == @HEAD_DIM@ && params.dv == @HEAD_DIM@) { + using AppendTileShapeQK = + cute::Shape, cute::Int<@APPENDKV_Q256_KV@>, cute::_32>; + using AppendTileShapePV = + cute::Shape, cute::_32, cute::Int<@APPENDKV_Q256_KV@>>; + using AppendTileShapeOutput = cute::Shape, cute::Int<@TILED_OUT@>>; + using AppendSubgroupLayoutQK = cute::Layout< + cute::Shape, cute::_1, cute::_1>, + cute::Stride>; + constexpr auto append_mode = @HEAD_DIM@ == 64 ? cutlass::fmha::collective::AppendKVMode::kStoreWide + : cutlass::fmha::collective::AppendKVMode::kStoreAndDirectLoad; + FMHAConfig::template run_paged(params); + return; + } +#endif + +#if @HEAD_DIM@ == 96 + if (params.b >= 8 && params.h_k >= 8) { + using DenseTileShapeQK = cute::Shape, cute::Int<64>, cute::_32>; + using DenseTileShapePV = cute::Shape, cute::_32, cute::Int<64>>; + using DenseTileShapeOutput = cute::Shape, cute::Int<96>>; + using DenseSubgroupLayoutQK = + cute::Layout, cute::_1, cute::_1>, cute::Stride>; + TORCH_CHECK(!use_sink, "sink is only supported for head_size == 64, got @HEAD_DIM@"); + run_paged_tile(params); + return; + } +#endif + +#if @HEAD_DIM@ == 256 && @HD256_PAGED_LARGE_TILE@ + if (params.seqlen_q >= @HD256_PAGED_LARGE_TILE_MIN_Q@ && + (params.total_kvnew == 0 || use_append_kv_prepass(params)) && + params.d == 256 && params.dv == 256) { + using LargeTileShapeQK = + cute::Shape, cute::Int<@HD256_PAGED_LARGE_TILE_KV@>, cute::_32>; + using LargeTileShapePV = + cute::Shape, cute::_32, cute::Int<@HD256_PAGED_LARGE_TILE_KV@>>; + using LargeTileShapeOutput = cute::Shape, cute::Int<256>>; + using LargeSubgroupLayoutQK = cute::Layout< + cute::Shape, cute::_1, cute::_1>, + cute::Stride>; + TORCH_CHECK(!use_sink, "sink is only supported for head_size == 64, got @HEAD_DIM@"); + AT_DISPATCH_BOOL_NO_RETURN(params.is_local, LocalMask, { + AT_DISPATCH_BOOL_NO_RETURN(params.is_causal, Causal, { + FMHAConfig::run_paged(params); + }); + }); + return; + } +#endif + +#if @HEAD_DIM@ == 64 + if (params.seqlen_q <= 256 && + static_cast(params.total_q) * 100 <= static_cast(params.b) * params.seqlen_q * 15 && + params.d == 64 && params.dv == 64) { + using SparseTileShapeQK = cute::Shape, cute::Int<32>, cute::_32>; + using SparseTileShapePV = cute::Shape, cute::_32, cute::Int<32>>; + using SparseTileShapeOutput = cute::Shape, cute::Int<64>>; + using SparseSubgroupLayoutQK = cute:: + Layout, cute::_1, cute::_1>, cute::Stride>; + AT_DISPATCH_BOOL_NO_RETURN(use_sink, Sink, { + AT_DISPATCH_BOOL_NO_RETURN(params.is_local, LocalMask, { + AT_DISPATCH_BOOL_NO_RETURN(params.is_causal, Causal, { + FMHAConfig::run_paged(params); + }); + }); + }); + return; + } + + if (params.seqlen_q <= 32 && params.d == 64 && params.dv == 64) { + using SmallTileShapeQK = cute::Shape, cute::Int<64>, cute::_32>; + using SmallTileShapePV = cute::Shape, cute::_32, cute::Int<64>>; + using SmallTileShapeOutput = cute::Shape, cute::Int<64>>; + using SmallSubgroupLayoutQK = cute:: + Layout, cute::_1, cute::_1>, cute::Stride>; + AT_DISPATCH_BOOL_NO_RETURN(use_sink, Sink, { + AT_DISPATCH_BOOL_NO_RETURN(params.is_local, LocalMask, { + AT_DISPATCH_BOOL_NO_RETURN(params.is_causal, Causal, { + FMHAConfig::run_paged(params); + }); + }); + }); + return; + } + + if (params.seqlen_q <= 64 && params.d == 64 && params.dv == 64) { + using SmallTileShapeQK = cute::Shape, cute::Int<64>, cute::_32>; + using SmallTileShapePV = cute::Shape, cute::_32, cute::Int<64>>; + using SmallTileShapeOutput = cute::Shape, cute::Int<64>>; + using SmallSubgroupLayoutQK = cute:: + Layout, cute::_1, cute::_1>, cute::Stride>; + AT_DISPATCH_BOOL_NO_RETURN(use_sink, Sink, { + AT_DISPATCH_BOOL_NO_RETURN(params.is_local, LocalMask, { + AT_DISPATCH_BOOL_NO_RETURN(params.is_causal, Causal, { + FMHAConfig::run_paged(params); + }); + }); + }); + return; + } + +#endif + +#if @HEAD_DIM@ == 128 + if (params.seqlen_q <= 256 && + static_cast(params.total_q) * 2 <= static_cast(params.b) * params.seqlen_q && + params.d == 128 && params.dv == 128) { + using MixedTileShapeQK = cute::Shape, cute::Int<64>, cute::_32>; + using MixedTileShapePV = cute::Shape, cute::_32, cute::Int<64>>; + using MixedTileShapeOutput = cute::Shape, cute::Int<128>>; + using MixedSubgroupLayoutQK = cute:: + Layout, cute::_1, cute::_1>, cute::Stride>; + TORCH_CHECK(!use_sink, "sink is only supported for head_size == 64, got @HEAD_DIM@"); + AT_DISPATCH_BOOL_NO_RETURN(params.is_local, LocalMask, { + AT_DISPATCH_BOOL_NO_RETURN(params.is_causal, Causal, { + FMHAConfig::run_paged(params); + }); + }); + return; + } + + if (params.seqlen_q <= 32 && params.d == 128 && params.dv == 128) { + using SmallTileShapeQK = cute::Shape, cute::Int<64>, cute::_32>; + using SmallTileShapePV = cute::Shape, cute::_32, cute::Int<64>>; + using SmallTileShapeOutput = cute::Shape, cute::Int<128>>; + using SmallSubgroupLayoutQK = cute:: + Layout, cute::_1, cute::_1>, cute::Stride>; + TORCH_CHECK(!use_sink, "sink is only supported for head_size == 64, got @HEAD_DIM@"); + AT_DISPATCH_BOOL_NO_RETURN(params.is_local, LocalMask, { + AT_DISPATCH_BOOL_NO_RETURN(params.is_causal, Causal, { + FMHAConfig::run_paged(params); + }); + }); + return; + } + + if (params.seqlen_q <= 64 && params.d == 128 && params.dv == 128) { + using SmallTileShapeQK = cute::Shape, cute::Int<64>, cute::_32>; + using SmallTileShapePV = cute::Shape, cute::_32, cute::Int<64>>; + using SmallTileShapeOutput = cute::Shape, cute::Int<128>>; + using SmallSubgroupLayoutQK = cute:: + Layout, cute::_1, cute::_1>, cute::Stride>; + TORCH_CHECK(!use_sink, "sink is only supported for head_size == 64, got @HEAD_DIM@"); + AT_DISPATCH_BOOL_NO_RETURN(params.is_local, LocalMask, { + AT_DISPATCH_BOOL_NO_RETURN(params.is_causal, Causal, { + FMHAConfig::run_paged(params); + }); + }); + return; + } +#endif + +#if @HEAD_DIM@ == 128 && @HD128_PAGED_LARGE_TILE@ + if (params.seqlen_q >= @HD128_PAGED_LARGE_TILE_MIN_Q@ && params.d == 128 && params.dv == 128) { + using LargeTileShapeQK = + cute::Shape, cute::Int<@HD128_PAGED_LARGE_TILE_KV@>, cute::_32>; + using LargeTileShapePV = + cute::Shape, cute::_32, cute::Int<@HD128_PAGED_LARGE_TILE_KV@>>; + using LargeTileShapeOutput = cute::Shape, cute::Int<128>>; + using LargeSubgroupLayoutQK = cute::Layout< + cute::Shape, cute::_1, cute::_1>, + cute::Stride>; + TORCH_CHECK(!use_sink, "sink is only supported for head_size == 64, got @HEAD_DIM@"); + AT_DISPATCH_BOOL_NO_RETURN(params.is_local, LocalMask, { + AT_DISPATCH_BOOL_NO_RETURN(params.is_causal, Causal, { + FMHAConfig::run_paged(params); + }); + }); + return; + } +#endif + #if @HEAD_DIM@ == 64 AT_DISPATCH_BOOL_NO_RETURN(use_sink, Sink, { AT_DISPATCH_BOOL_NO_RETURN(params.is_local, LocalMask, { diff --git a/src/sycl/kernels/flash_attention_v2/xe_fmha_fwd_prefill_nopage_kernel.cpp.in b/src/sycl/kernels/flash_attention_v2/xe_fmha_fwd_prefill_nopage_kernel.cpp.in index 01940f6e7..7c418398c 100644 --- a/src/sycl/kernels/flash_attention_v2/xe_fmha_fwd_prefill_nopage_kernel.cpp.in +++ b/src/sycl/kernels/flash_attention_v2/xe_fmha_fwd_prefill_nopage_kernel.cpp.in @@ -42,19 +42,10 @@ namespace prefill { -template <> -void FmhaPrefillNpRunner<@HEAD_DIM@>::operator()(const Arguments& params) const { - // Non-paged (contiguous ragged) KV cache. Uses an independent tile-shape - // configuration so it can be tuned separately from the paged path. - TORCH_CHECK( - params.softmax_sink_ptr == nullptr, - "No-page attention does not support sink logits (softmax_sink_ptr must be null)"); - using TileShapeQK = cute::Shape, cute::Int<@TILED_KV_NP@>, cute::_32>; - using TileShapePV = cute::Shape, cute::_32, cute::Int<@TILED_KV_NP@>>; - using TileShapeOutput = cute::Shape, cute::Int<((@HEAD_DIM@ + 31) / 32) * 32>>; - using SubgroupLayoutQK = cute:: - Layout, cute::_1, cute::_1>, cute::Stride>; - // No sink support for no_page path; always use Sink=false +namespace { + +template +void run_nopaged_tile(const Arguments& params) { AT_DISPATCH_BOOL_NO_RETURN(params.is_local, LocalMask, { AT_DISPATCH_BOOL_NO_RETURN(params.is_causal, Causal, { using Element = bfloat16_t; @@ -77,6 +68,52 @@ void FmhaPrefillNpRunner<@HEAD_DIM@>::operator()(const Arguments& params) const }); } +} // namespace + +template <> +void FmhaPrefillNpRunner<@HEAD_DIM@>::operator()(const Arguments& params) const { + TORCH_CHECK( + params.softmax_sink_ptr == nullptr, + "No-page attention does not support sink logits (softmax_sink_ptr must be null)"); + using TileShapeQK = cute::Shape, cute::Int<@TILED_KV_NP@>, cute::_32>; + using TileShapePV = cute::Shape, cute::_32, cute::Int<@TILED_KV_NP@>>; + using TileShapeOutput = cute::Shape, cute::Int<((@HEAD_DIM@ + 31) / 32) * 32>>; + using SubgroupLayoutQK = cute:: + Layout, cute::_1, cute::_1>, cute::Stride>; + +#if @HEAD_DIM@ == 96 && @HD96_NP_SMALL_TILE@ + if (params.seqlen_q <= @HD96_NP_SMALL_TILE_MAX_Q@) { + using SmallTileShapeQK = + cute::Shape, cute::Int<@HD96_NP_SMALL_TILE_KV@>, cute::_32>; + using SmallTileShapePV = + cute::Shape, cute::_32, cute::Int<@HD96_NP_SMALL_TILE_KV@>>; + using SmallTileShapeOutput = cute::Shape, cute::Int<96>>; + using SmallSubgroupLayoutQK = cute::Layout< + cute::Shape, cute::_1, cute::_1>, + cute::Stride>; + run_nopaged_tile(params); + return; + } +#endif + +#if @HEAD_DIM@ == 96 && @HD96_NP_LARGE_TILE@ + if (params.seqlen_q >= @HD96_NP_LARGE_TILE_MIN_Q@) { + using LargeTileShapeQK = + cute::Shape, cute::Int<@HD96_NP_LARGE_TILE_KV@>, cute::_32>; + using LargeTileShapePV = + cute::Shape, cute::_32, cute::Int<@HD96_NP_LARGE_TILE_KV@>>; + using LargeTileShapeOutput = cute::Shape, cute::Int<96>>; + using LargeSubgroupLayoutQK = cute::Layout< + cute::Shape, cute::_1, cute::_1>, + cute::Stride>; + run_nopaged_tile(params); + return; + } +#endif + + run_nopaged_tile(params); +} + template struct FmhaPrefillNpRunner<@HEAD_DIM@>; } // namespace prefill diff --git a/src/sycl/kernels/flash_attention_v2/xe_fmha_fwd_prefill_runner.hpp b/src/sycl/kernels/flash_attention_v2/xe_fmha_fwd_prefill_runner.hpp index 8789e8dc3..3aae33cff 100644 --- a/src/sycl/kernels/flash_attention_v2/xe_fmha_fwd_prefill_runner.hpp +++ b/src/sycl/kernels/flash_attention_v2/xe_fmha_fwd_prefill_runner.hpp @@ -43,6 +43,7 @@ #include "sycl/Utils.h" #include "sycl/comm/common.h" #include "sycl/kernels/flash_attention_v2/collective/fmha_fusion.hpp" +#include "sycl/kernels/flash_attention_v2/kernel/xe_fmha_append_kv.hpp" #include "sycl/kernels/flash_attention_v2/kernel/xe_fmha_fwd_kernel.hpp" #include "sycl/kernels/flash_attention_v2/kernel/xe_tile_scheduler.hpp" @@ -84,9 +85,9 @@ struct Arguments { void* __restrict__ softmax_lseaccum_ptr; // The dimensions. - int b, seqlen_q, seqlen_k, seqlen_knew, d, d_rounded, rotary_dim; - int total_q, total_k; - int total_knew = 0; + int b, seqlen_q, seqlen_kvcache, seqlen_kvnew, d, d_rounded, rotary_dim; + int total_q, total_kvcache; + int total_kvnew = 0; int b_k; // When having KV cache and with cache_batch_idx, K & V might have larger batch size than Q int dv, dv_rounded; // For the case where V headdim is different from Q/K headdim @@ -104,12 +105,13 @@ struct Arguments { // array of length b+1 holding starting offset of each sequence. int* __restrict__ cu_seqlens_q; int* __restrict__ cu_seqlens_k; - int* __restrict__ cu_seqlens_knew; + int* __restrict__ cu_seqlens_kvnew; int* __restrict__ leftpad_k; // If provided, the actual length of each q/k sequence. int* __restrict__ seqused_q; int* __restrict__ seqused_k; + int* __restrict__ kv_cache_seqlens; // The stride between rows of Oaccum. int64_t oaccum_split_stride; @@ -123,16 +125,16 @@ struct Arguments { int64_t lseaccum_head_stride; // The K_new and V_new matrices. - void* __restrict__ knew_ptr; - void* __restrict__ vnew_ptr; + void* __restrict__ k_new_ptr; + void* __restrict__ v_new_ptr; // The stride between rows of the Q, K and V matrices. - int64_t knew_batch_stride; - int64_t vnew_batch_stride; - int64_t knew_row_stride; - int64_t vnew_row_stride; - int64_t knew_head_stride; - int64_t vnew_head_stride; + int64_t k_new_batch_stride; + int64_t v_new_batch_stride; + int64_t k_new_row_stride; + int64_t v_new_row_stride; + int64_t k_new_head_stride; + int64_t v_new_head_stride; void* __restrict__ qv_ptr; int64_t qv_batch_stride; @@ -192,6 +194,38 @@ using LayoutK = cutlass::layout::ColumnMajor; using LayoutV = cutlass::layout::RowMajor; using LayoutO = cutlass::layout::RowMajor; +constexpr int kAppendKVPrepassMinTokens = 64; + +inline bool use_append_kv_prepass(const Arguments& params) { + return params.total_kvnew >= kAppendKVPrepassMinTokens && + params.k_new_ptr != nullptr && params.v_new_ptr != nullptr && + params.kv_cache_seqlens != nullptr && !params.is_e4m3 && !params.is_e5m2; +} + +template +void launch_append_kv_prepass(const Arguments& params) { + using AppendKernel = cutlass::fmha::kernel::XeFMHAAppendKVKernel; + typename AppendKernel::Arguments args{}; + args.ptr_K_cache = static_cast(params.k_ptr); + args.ptr_V_cache = static_cast(params.v_ptr); + args.ptr_K_new = static_cast(params.k_new_ptr); + args.ptr_V_new = static_cast(params.v_new_ptr); + args.ptr_cu_seqlens_k_new = params.cu_seqlens_kvnew; + args.seq_len_kv_new = params.seqlen_kvnew; + args.ptr_cache_seqlens = params.kv_cache_seqlens; + args.ptr_page_table = params.page_table; + args.page_size = params.page_size; + args.max_num_pages_per_seq = params.max_num_pages_per_seq; + args.batch = params.b; + args.num_heads_kv = params.h_k; + args.head_size_qk = params.d; + args.head_size_vo = params.dv; + args.max_seq_len_kv_new = params.seqlen_kvnew; + args.skip_batch_mask = static_cast(params.skip_batch_mask_ptr); + TORCH_CHECK(AppendKernel::can_implement(args), "AppendKV pre-pass cannot implement these arguments"); + launch(AppendKernel::to_underlying_arguments(args, nullptr)); +} + template struct PrefillRunner { using StrideQ = typename FMHAPrefillKernel::StrideQ; @@ -231,16 +265,16 @@ struct PrefillRunner { get<0>(problem_size_for_init) = 1; // concentrated batch get<1>(problem_size_for_init) = params.h; get<3>(problem_size_for_init) = params.total_q; - get<4>(problem_size_for_init) = params.total_knew; - get<5>(problem_size_for_init) = params.total_k; + get<4>(problem_size_for_init) = params.total_kvnew; + get<5>(problem_size_for_init) = params.total_kvcache; ProblemShapeType problem_size_for_launch{ .batch = get<0>(problem_size), .num_heads_q = get<1>(problem_size), .num_heads_kv = get<2>(problem_size), .seq_len_qo = {params.seqlen_q, params.total_q, nullptr}, - .seq_len_kv = {params.seqlen_knew, params.total_knew}, - .seq_len_kv_cache = {params.seqlen_k, params.total_k}, + .seq_len_kv = {params.seqlen_kvnew, params.total_kvnew}, + .seq_len_kv_cache = {params.seqlen_kvcache, params.total_kvcache}, .head_size_qk = get<6>(problem_size), .head_size_vo = get<7>(problem_size), }; @@ -251,7 +285,14 @@ struct PrefillRunner { /// Initialize operands to be used in the GEMM and reference GEMM ProblemShapeType initialize(const Arguments& params) { auto problem_shape_in = cute::make_tuple( - params.b, params.h, params.h_k, params.seqlen_q, params.seqlen_knew, params.seqlen_k, params.d, params.dv); + params.b, + params.h, + params.h_k, + params.seqlen_q, + params.seqlen_kvnew, + params.seqlen_kvcache, + params.d, + params.dv); ProblemShapeType shape; decltype(problem_shape_in) problem_size; @@ -283,7 +324,7 @@ struct PrefillRunner { if constexpr (isVarLen) { shape.seq_len_qo.cumulative_length = params.cu_seqlens_q; - shape.seq_len_kv.cumulative_length = params.cu_seqlens_knew; + shape.seq_len_kv.cumulative_length = params.cu_seqlens_kvnew; shape.seq_len_kv_cache.cumulative_length = params.cu_seqlens_k; } @@ -293,14 +334,30 @@ struct PrefillRunner { cutlass::Status run(const Arguments& params, const cutlass::KernelHardwareInfo& hw_info) { ProblemShapeType shape = initialize(params); + typename FMHAPrefillKernel::MainloopArguments mainloop_args{ + params.softmax_scale, + params.page_table, + params.page_size, + params.max_num_pages_per_seq, + params.window_size_left, + params.window_size_right}; + if constexpr (CollectiveMainloop::AppendKV) { + mainloop_args.ptr_K_new = static_cast(params.k_new_ptr); + mainloop_args.ptr_V_new = static_cast(params.v_new_ptr); + mainloop_args.ptr_cu_seqlens_k_new = params.cu_seqlens_kvnew; + mainloop_args.ptr_kv_cache_seqlens = params.kv_cache_seqlens; + mainloop_args.seq_len_kv_new = params.seqlen_kvnew; + mainloop_args.total_k_new = params.total_kvnew; + } + typename FMHAPrefillKernel::Arguments arguments{ { shape, static_cast(params.q_ptr), stride_Q, - nullptr, + CollectiveMainloop::AppendKV ? static_cast(params.k_new_ptr) : nullptr, stride_K, - nullptr, + CollectiveMainloop::AppendKV ? static_cast(params.v_new_ptr) : nullptr, stride_V, static_cast(params.o_ptr), stride_O, @@ -313,14 +370,7 @@ struct PrefillRunner { params.k_scale_ptr, params.v_scale_ptr, }, - { - params.softmax_scale, - params.page_table, - params.page_size, - params.max_num_pages_per_seq, - params.window_size_left, - params.window_size_right, - }, + mainloop_args, {}, hw_info}; @@ -381,7 +431,12 @@ struct FMHAConfig { decltype(cutlass::fmha::collective::get_sg_layout_pv(SubgroupLayoutQK{})), SubgroupLayoutPV_>; - template + template < + bool isVarLen, + bool CachedKV, + bool PagedKV, + cutlass::fmha::collective::AppendKVMode AppendMode, + class Scheduler> static int run(const Arguments& params) { // The KernelHardwareInfo struct holds the number of EUs on the GPU with a given device ID. This // information is used by the underlying kernel. @@ -431,7 +486,9 @@ struct FMHAConfig { GmemTiledCopyV, GmemTiledCopyK_cache, GmemTiledCopyV_cache, - LocalMask>; + LocalMask, + false, + AppendMode>; // Epilogue using CollectiveEpilogue = @@ -459,15 +516,70 @@ struct FMHAConfig { } // Paged KV cache: the page table encodes absolute KV positions. + template < + cutlass::fmha::collective::AppendKVMode AppendMode = int(get<1>(TileShapeOutput{})) == 64 + ? cutlass::fmha::collective::AppendKVMode::kStoreWide + : cutlass::fmha::collective::AppendKVMode::kStore> static int run_paged(const Arguments& params) { - // template - return run(params); + TORCH_CHECK(params.cu_seqlens_q != nullptr, "paged prefill requires cu_seqlens_q"); + TORCH_CHECK(params.cu_seqlens_k != nullptr, "paged prefill requires per-batch cache lengths in cu_seqlens_k"); + TORCH_CHECK(params.page_table != nullptr, "paged prefill requires page_table"); + TORCH_CHECK(params.page_size > 0, "paged prefill requires a positive page_size"); + TORCH_CHECK(params.max_num_pages_per_seq > 0, "paged prefill requires max_num_pages_per_seq"); + TORCH_CHECK( + params.seqlen_q > 0 && params.seqlen_kvcache > 0, "paged prefill requires positive max sequence lengths"); + TORCH_CHECK( + params.total_q > 0 && params.total_kvcache > 0, "paged prefill requires positive total sequence lengths"); + bool const has_append = params.total_kvnew > 0 && params.k_new_ptr != nullptr && params.v_new_ptr != nullptr && + params.kv_cache_seqlens != nullptr; + if (has_append) { + if (use_append_kv_prepass(params)) { + launch_append_kv_prepass(params); + // Re-enter the existing no-append dispatch rather than instantiating a + // second copy of every attention tile for this runtime-only decision. + Arguments cached_kv_params = params; + cached_kv_params.k_new_ptr = nullptr; + cached_kv_params.v_new_ptr = nullptr; + cached_kv_params.cu_seqlens_kvnew = nullptr; + cached_kv_params.kv_cache_seqlens = nullptr; + cached_kv_params.total_kvnew = 0; + return run_paged(cached_kv_params); + } + return run(params); + } + return run< + true, + true, + true, + cutlass::fmha::collective::AppendKVMode::kNone, + cutlass::fmha::kernel::XeFHMAIndividualTileScheduler>(params); } // Non-paged (contiguous ragged) KV cache: addressed via cu_seqlens_k offsets. static int run_nopaged(const Arguments& params) { - // template - return run(params); + TORCH_CHECK(params.cu_seqlens_q != nullptr, "non-paged prefill requires cu_seqlens_q"); + TORCH_CHECK(params.cu_seqlens_k != nullptr, "non-paged prefill requires cumulative cu_seqlens_k"); + TORCH_CHECK(params.page_table == nullptr, "non-paged prefill expects page_table to be null"); + TORCH_CHECK( + params.seqlen_q > 0 && params.seqlen_kvcache > 0, "non-paged prefill requires positive max sequence lengths"); + TORCH_CHECK( + params.total_q > 0 && params.total_kvcache > 0, "non-paged prefill requires positive total sequence lengths"); + bool const has_append = params.total_kvnew > 0 && params.k_new_ptr != nullptr && params.v_new_ptr != nullptr && + params.kv_cache_seqlens != nullptr; + if (has_append) { + return run< + true, + true, + false, + cutlass::fmha::collective::AppendKVMode::kStore, + cutlass::fmha::kernel::XeFHMAIndividualTileScheduler>(params); + } + return run< + true, + true, + false, + cutlass::fmha::collective::AppendKVMode::kNone, + cutlass::fmha::kernel::XeFHMAIndividualTileScheduler>(params); } static int run(const Arguments& params) { diff --git a/src/torch_extension_sycl.cc b/src/torch_extension_sycl.cc index 30e594eb8..ba0668ef4 100644 --- a/src/torch_extension_sycl.cc +++ b/src/torch_extension_sycl.cc @@ -175,7 +175,10 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { " int num_kv_splits," " bool? pack_gqa," " int sm_margin," - " Tensor(a!)? out=None) -> (Tensor(a!), Tensor, Tensor, Tensor)"); + " Tensor(a!)? out=None," + " Tensor? k_new=None," + " Tensor? v_new=None," + " Tensor? cu_seqlens_k_new=None) -> (Tensor(a!), Tensor, Tensor, Tensor)"); m.impl("fwd", torch::kXPU, make_pytorch_shim(&mha_fwd)); #endif // USE_FMHA diff --git a/tests/test_flash_attention.py b/tests/test_flash_attention.py index f78705a8b..08947b329 100644 --- a/tests/test_flash_attention.py +++ b/tests/test_flash_attention.py @@ -59,7 +59,7 @@ def is_fa3_supported(device=None) -> bool: DISABLE_SPLIT = True DISABLE_PAGEDKV = False -DISABLE_APPENDKV = True +DISABLE_APPENDKV = False DISABLE_LOCAL = True DISABLE_SOFTCAP = True DISABLE_PACKGQA = True @@ -479,7 +479,7 @@ def generate_qkv( "dtype", [torch.bfloat16] + ([torch.float8_e4m3fn] if not DISABLE_FP8 else []) ) @pytest.mark.parametrize("nheads_q,nheads_kv", [(16, 16), (16, 4)]) -@pytest.mark.parametrize("new_kv", [False]) +@pytest.mark.parametrize("new_kv", [False, True]) @pytest.mark.parametrize("causal,local", [(False, True), (False, False), (True, False)]) @pytest.mark.parametrize("use_sinks", [True, False]) @pytest.mark.parametrize("seqlen_new_eq_seqlen_q", [True]) @@ -691,10 +691,13 @@ def test_flash_attn_kvcache( dtype, dtype_ref, ) + cache_seqlens_high = seqlen_k - seqlen_new + 1 if new_kv else seqlen_k + if cache_seqlens_high <= seqlen_q: + pytest.skip("new_kv requires room for appended KV") cache_seqlens = torch.randint( seqlen_q, # If we don't use seqlen_q in the case of causal and rotary, cos/sin won't be long enough - seqlen_k, + cache_seqlens_high, (batch_size,), dtype=torch.int32, device=device, @@ -880,6 +883,7 @@ def test_flash_attn_kvcache( cu_seqlens_q=cu_seqlens_q, cu_seqlens_k_new=cu_seqlens_k_new, max_seqlen_q=max_seqlen_q, + max_seqlen_k=seqlen_new if new_kv else max_seqlen_k, rotary_seqlens=rotary_seqlens, causal=causal, window_size=window_size,