From 06cbec02074ceaae48a2bddcf2d29e60f54e0d17 Mon Sep 17 00:00:00 2001 From: Yuankun Shi Date: Wed, 15 Jul 2026 23:24:12 -0700 Subject: [PATCH 01/16] Replace chunk prefill with AppendKV FMHA prefill --- include/sgl_flash_kernel_ops.h | 34 +++ python/sgl_kernel/flash_attn.py | 104 +++++-- src/FMHAPrefillXe20.cmake | 19 +- src/sycl/flash_attention.cpp | 226 +++++++++++++- .../collective/xe_fmha_fwd_mainloop.hpp | 275 ++++++++++++++++-- .../kernel/xe_fmha_fwd_kernel.hpp | 55 +++- .../xe_fmha_fwd_prefill_kernel.cpp.in | 73 +++++ .../xe_fmha_fwd_prefill_runner.hpp | 62 +++- src/torch_extension_sycl.cc | 34 +++ tests/test_flash_attention.py | 10 +- 10 files changed, 802 insertions(+), 90 deletions(-) diff --git a/include/sgl_flash_kernel_ops.h b/include/sgl_flash_kernel_ops.h index 44b9e61a4..16842c400 100644 --- a/include/sgl_flash_kernel_ops.h +++ b/include/sgl_flash_kernel_ops.h @@ -76,6 +76,40 @@ std::tuple mha_fwd( int const sm_margin, std::optional& out_); +std::tuple mha_fwd_appendkv( + const at::Tensor& q, + const at::Tensor& k, + const at::Tensor& v, + std::optional& q_v_, + const at::Tensor& cu_seqlens_q, + const at::Tensor& cu_seqlens_k, + int max_seqlen_q, + int max_seqlen_k, + std::optional& page_table, + std::optional& kv_batch_idx_, + std::optional& leftpad_k_, + std::optional& rotary_cos_, + std::optional& rotary_sin_, + std::optional& seqlens_rotary_, + std::optional& q_descale_, + std::optional& k_descale_, + std::optional& v_descale_, + float const softmax_scale, + std::optional& sinks, + bool is_causal, + int window_size_left, + int window_size_right, + float const softcap, + bool const is_rotary_interleaved, + std::optional& scheduler_metadata_, + int num_kv_splits, + std::optional pack_gqa_, + int const sm_margin, + std::optional& out_, + std::optional& k_new_, + std::optional& v_new_, + std::optional& cu_seqlens_k_new_); + void flash_mla_decode( torch::Tensor& out, const torch::Tensor& q_nope, diff --git a/python/sgl_kernel/flash_attn.py b/python/sgl_kernel/flash_attn.py index 78740c5c8..bf2607af9 100644 --- a/python/sgl_kernel/flash_attn.py +++ b/python/sgl_kernel/flash_attn.py @@ -264,37 +264,79 @@ 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 - out, softmax_lse, *rest = torch.ops.sgl_kernel.fwd.default( - q, - k_cache, - v_cache, - qv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - 1, - page_table, - cache_batch_idx, - cache_leftpad, - rotary_cos, - rotary_sin, - rotary_seqlens, - q_descale, - k_descale, - v_descale, - softmax_scale, - sinks, - causal, - window_size[0], - window_size[1], - softcap, - rotary_interleaved, - scheduler_metadata, - num_splits, - pack_gqa, - sm_margin, - out, - ) + has_new_kv = k is not None or v is not None or cu_seqlens_k_new is not None + 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_appendkv.default( + q, + k_cache, + v_cache, + qv, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + native_max_seqlen_k, + page_table, + cache_batch_idx, + cache_leftpad, + rotary_cos, + rotary_sin, + rotary_seqlens, + q_descale, + k_descale, + v_descale, + softmax_scale, + sinks, + causal, + window_size[0], + window_size[1], + softcap, + rotary_interleaved, + scheduler_metadata, + num_splits, + pack_gqa, + sm_margin, + out, + k, + v, + cu_seqlens_k_new, + ) + else: + out, softmax_lse, *rest = torch.ops.sgl_kernel.fwd.default( + q, + k_cache, + v_cache, + qv, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + 1, + page_table, + cache_batch_idx, + cache_leftpad, + rotary_cos, + rotary_sin, + rotary_seqlens, + q_descale, + k_descale, + v_descale, + softmax_scale, + sinks, + causal, + window_size[0], + window_size[1], + softcap, + rotary_interleaved, + scheduler_metadata, + num_splits, + pack_gqa, + sm_margin, + out, + ) return (out, softmax_lse, *rest) if return_softmax_lse else out diff --git a/src/FMHAPrefillXe20.cmake b/src/FMHAPrefillXe20.cmake index 9e96270dc..bdd1c4920 100644 --- a/src/FMHAPrefillXe20.cmake +++ b/src/FMHAPrefillXe20.cmake @@ -33,9 +33,14 @@ set(FMHA_PREFILL_TILED_Q_96 128) set(FMHA_PREFILL_TILED_KV_96 64) set(FMHA_PREFILL_NUM_SG_96 8) -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_KV_192 64) @@ -96,6 +101,16 @@ 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}) + 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) diff --git a/src/sycl/flash_attention.cpp b/src/sycl/flash_attention.cpp index e45e238c6..155530f09 100644 --- a/src/sycl/flash_attention.cpp +++ b/src/sycl/flash_attention.cpp @@ -185,6 +185,8 @@ std::vector mha_fwd_nopage( 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). + params.knew_ptr = nullptr; + params.vnew_ptr = nullptr; params.cu_seqlens_knew = nullptr; params.seqlen_knew = 0; params.total_knew = 0; @@ -213,6 +215,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 +521,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 +781,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.knew_ptr = nullptr; + params.vnew_ptr = nullptr; + params.cu_seqlens_knew = nullptr; + params.cache_seqlens_old = nullptr; + params.seqlen_knew = 0; + params.total_knew = 0; params.softmax_lse_ptr = softmax_lse.data_ptr(); @@ -855,7 +870,7 @@ std::vector mha_fwd_nopage( return {out, softmax_lse, out_accum, softmax_lse_accum}; } -std::vector mha_fwd( +std::vector mha_fwd_appendkv( const at::Tensor& q, // (b, s_q, h, d) or (total_q, h, d) if there is cu_seqlens_q const at::Tensor& k, // (b_k, s_k, h_k, d) or (total_k, h_k, d) if there is cu_seqlens_k or (num_pages, page_size, // h_k, d) if there is page_table. @@ -889,7 +904,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 +933,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, @@ -1101,6 +1122,53 @@ std::vector mha_fwd( params.page_size = page_size; params.num_pages = num_pages; + params.knew_ptr = nullptr; + params.vnew_ptr = nullptr; + params.cu_seqlens_knew = nullptr; + params.cache_seqlens_old = nullptr; + params.seqlen_knew = 0; + params.total_knew = 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_knew = 0; + int seqlen_knew = max_seqlen_k > 0 ? max_seqlen_k : max_seqlen_q; + if (k_new.dim() == 3) { + total_knew = k_new.size(0); + CHECK_SHAPE(k_new, total_knew, num_heads_k, head_size); + CHECK_SHAPE(v_new, total_knew, num_heads_k, head_size_v); + TORCH_CHECK( + cu_seqlens_k_new_.has_value() || seqlen_knew > 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_knew = batch_size * k_new_seqlen; + seqlen_knew = 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_knew = cu_seqlens_k_new.data_ptr(); + } + params.knew_ptr = k_new.data_ptr(); + params.vnew_ptr = v_new.data_ptr(); + params.cache_seqlens_old = cu_seqlens_k.data_ptr(); + params.seqlen_knew = seqlen_knew; + params.total_knew = total_knew; + } + if (q_v_.has_value()) { TORCH_CHECK(head_size <= 64, "q_v is only supported for head_size <= 64"); TORCH_CHECK( @@ -1191,6 +1259,75 @@ std::vector mha_fwd( return {out, softmax_lse, out_accum, softmax_lse_accum}; } +std::vector mha_fwd( + const at::Tensor& q, // (b, s_q, h, d) or (total_q, h, d) if there is cu_seqlens_q + const at::Tensor& k, // (b_k, s_k, h_k, d) or (total_k, h_k, d) if there is cu_seqlens_k or (num_pages, page_size, + // h_k, d) if there is page_table. + const at::Tensor& v, // (b_k, s_k, h_k, dv) or (total_k, h_k, dv) if there is cu_seqlens_k or (num_pages, + // page_size, h_k, dv) if there is page_table. + std::optional& q_v_, // (b, s_q, h, dv) or (total_q_new, h, dv) if there is cu_seqlens_q + const at::Tensor& cu_seqlens_q, // b+1 + const at::Tensor& cu_seqlens_k, // b+1 + int max_seqlen_q, + int max_seqlen_k, + std::optional& page_table, // (b_k, max_num_pages_per_seq) + std::optional& kv_batch_idx_, // b. indices to index into the KV cache + std::optional& leftpad_k_, // b + std::optional& rotary_cos_, // seqlen_ro x (rotary_dim / 2) + std::optional& rotary_sin_, // seqlen_ro x (rotary_dim / 2) + std::optional& seqlens_rotary_, // b + std::optional& q_descale_, // (b, h_k), not (b, h) + std::optional& k_descale_, // (b, h_k) + std::optional& v_descale_, // (b, h_k) + const float softmax_scale_, + std::optional& sinks_, + bool is_causal, + int window_size_left, + int window_size_right, + float const softcap, + bool const is_rotary_interleaved, // if true, rotary combines indices 0 & 1, else indices 0 & rotary_dim / 2 + std::optional& scheduler_metadata_, // (b + 1) + int num_splits, + std::optional pack_gqa_, + int const sm_margin, + std::optional out_opt = std::nullopt, + std::optional skip_batch_mask_opt = std::nullopt) { + return mha_fwd_appendkv( + 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_splits, + pack_gqa_, + sm_margin, + std::move(out_opt), + std::move(skip_batch_mask_opt), + std::nullopt, + std::nullopt, + std::nullopt); +} + } // namespace prefill namespace chunkprefill { @@ -1406,9 +1543,90 @@ std::tuple mha_fwd( // a single sequence with max_seqlen_q > 1 is prefill return dispatch(prefill::mha_fwd, std::nullopt); } else { - // Chunk prefill path + // Prefill path // Paged attn with max_seqlen_q > 1 and batch_size > 1 - return dispatch(chunkprefill::mha_fwd); + return dispatch(prefill::mha_fwd, std::nullopt); } } + +std::tuple mha_fwd_appendkv( + const at::Tensor& q, // (total_q, h, d) — ragged 3D + const at::Tensor& k, // (total_k, h_k, d) if non-paged, or (num_pages, page_size, h_k, d) if paged + const at::Tensor& v, // (total_k, h_k, dv) if non-paged, or (num_pages, page_size, h_k, dv) if paged + std::optional& q_v_, // (total_q, h, dv) — not yet supported + const at::Tensor& cu_seqlens_q, // b+1 + const at::Tensor& cu_seqlens_k, // b+1 + int max_seqlen_q, + int max_seqlen_k, + std::optional& page_table, // (b_k, max_num_pages_per_seq) + std::optional& kv_batch_idx_, // b. indices to index into the KV cache + std::optional& leftpad_k_, // b + std::optional& rotary_cos_, // seqlen_ro x (rotary_dim / 2) + std::optional& rotary_sin_, // seqlen_ro x (rotary_dim / 2) + std::optional& seqlens_rotary_, // b + std::optional& q_descale_, // (b, h_k), not (b, h) + std::optional& k_descale_, // (b, h_k) + std::optional& v_descale_, // (b, h_k) + const float softmax_scale_, + std::optional& sinks_, + bool is_causal, + int window_size_left, + int window_size_right, + float const softcap, + bool const is_rotary_interleaved, // if true, rotary combines indices 0 & 1, else indices 0 & rotary_dim / 2 + std::optional& scheduler_metadata_, // (b + 1) + int num_kv_splits, + std::optional pack_gqa_, + int const sm_margin, + 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)"); + if (out_.has_value()) { + const at::Tensor& out_val = out_.value(); + TORCH_CHECK(out_val.scalar_type() == q.scalar_type(), "out dtype must match q dtype"); + TORCH_CHECK( + out_val.dim() == 3 && out_val.size(0) == q.size(0) && out_val.size(1) == q.size(1) && + out_val.size(2) == v.size(-1), + "out shape must be [total_q, num_heads, head_size_v]"); + TORCH_CHECK(out_val.device() == q.device(), "out must be on the same device as q"); + TORCH_CHECK(out_val.stride(-1) == 1, "out must have a contiguous last dimension"); + } + auto to_tuple = [](std::vector v) { return std::make_tuple(v[0], v[1], v[2], v[3]); }; + return to_tuple(prefill::mha_fwd_appendkv( + 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_, + std::nullopt, + 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..713ec7956 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,21 @@ using namespace cute; ///////////////////////////////////////////////////////////////////////////////////////////////// +template +struct AppendKVParams {}; + +template +struct AppendKVParams { + ElementK_ const* ptr_K_new = nullptr; + ElementV_ const* ptr_V_new = nullptr; + int const* ptr_cu_seqlens_k_new = nullptr; + int const* ptr_cache_seqlens = nullptr; + int seq_len_kv_new = 0; + int total_k_new = 0; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + template < class DispatchPolicy_, bool CausalMask_, @@ -75,7 +93,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, + bool AppendKV_ = false> struct FMHAFwdMainloop { static_assert(cutlass::detail::dependent_false, "Could not find a mainloop specialization."); }; @@ -101,7 +120,8 @@ template < class TiledCopyK_cache_, class TiledCopyV_cache_, bool LocalMask_, - bool PackGQA_> + bool PackGQA_, + bool AppendKV_> struct FMHAFwdMainloop< XeDefault, CausalMask_, @@ -121,7 +141,8 @@ struct FMHAFwdMainloop< TiledCopyK_cache_, TiledCopyV_cache_, LocalMask_, - PackGQA_> { + PackGQA_, + AppendKV_> { // // Type Aliases // @@ -199,6 +220,9 @@ 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 bool AppendKV = AppendKV_; + using AppendKVStorage = + AppendKVParams; // User-facing arguments struct Arguments { @@ -208,6 +232,7 @@ struct FMHAFwdMainloop< int max_num_pages_per_seq = 0; int window_size_left = -1; int window_size_right = -1; + AppendKVStorage append{}; }; // Kernel-facing parameters @@ -233,7 +258,8 @@ struct FMHAFwdMainloop< args.page_size, args.max_num_pages_per_seq, args.window_size_left, - args.window_size_right}; + args.window_size_right, + args.append}; } CUTLASS_HOST_DEVICE static bool can_implement(Arguments const&) { @@ -254,6 +280,185 @@ struct FMHAFwdMainloop< return params.ptr_page_table[batch_offset + next_page_logical_idx] * tiles_per_page + K % tiles_per_page; } + CUTLASS_DEVICE + int get_k_new_len(int batch) const { + if constexpr (AppendKV) { + if (params.append.ptr_K_new == nullptr || params.append.ptr_V_new == nullptr || + params.append.ptr_cache_seqlens == nullptr || params.append.total_k_new <= 0 || + (params.append.ptr_cu_seqlens_k_new == nullptr && params.append.seq_len_kv_new <= 0)) { + return 0; + } + if (params.append.ptr_cu_seqlens_k_new != nullptr) { + return params.append.ptr_cu_seqlens_k_new[batch + 1] - params.append.ptr_cu_seqlens_k_new[batch]; + } + return params.append.seq_len_kv_new; + } else { + (void)batch; + return 0; + } + } + + 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) const { + if constexpr (AppendKV) { + int const new_len_total = get_k_new_len(batch); + int const new_len = + append_store_len < 0 ? new_len_total : (append_store_len < new_len_total ? append_store_len : new_len_total); + 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.append.ptr_cu_seqlens_k_new != nullptr ? params.append.ptr_cu_seqlens_k_new[batch] + : batch * params.append.seq_len_kv_new; + int const cache_len_old = params.append.ptr_cache_seqlens[batch]; + 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) { + // NHD bf16 rows are 16B-aligned when head_size is divisible by 8. + constexpr int kVecElems = 8; + using StoreVec = cutlass::ulonglong2; + if (head_size_qk == head_size_vo && (head_size_qk % kVecElems) == 0) { + // 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; + } + int const head_size = head_size_qk; + int const vecs_per_token = head_size / kVecElems; + 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; + + 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.append.ptr_K_new + src); + StoreVec const v_value = *reinterpret_cast(params.append.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; + } + } + + 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 dst_row = single_row_base + new_tok; + if constexpr (PagedKV) { + if (!single_dst_page) { + int const dst_tok = cache_len_old + new_tok; + 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; + } + } + + 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.append.ptr_K_new + src); + StoreVec const v_value = *reinterpret_cast(params.append.ptr_V_new + src); + *reinterpret_cast(&K_dst(dst_row, d)) = k_value; + *reinterpret_cast(&V_dst(d, dst_row)) = v_value; + } + } + 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.append.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.append.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; + } + } + template CUTLASS_DEVICE void operator()( TensorQ2D const& Q_2D, // (q,d) @@ -271,8 +476,11 @@ 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, 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,19 +573,27 @@ struct FMHAFwdMainloop< // Kernel // ------ + if constexpr (AppendKV) { + store_kv_new(K_cache_2D, V_cache_2D, l_coord, kv_head, num_heads_kv, thr_id, append_store_len); + barrier(); + } + /* 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); - } 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)); + bool has_prefetch_k = blk_k0 < blk_k1 && blk_k0 < kblocks_cache; + if (has_prefetch_k) { + if constexpr (PagedKV) { + next_page_idx = get_physical_k_tile(blk_k0, l_coord, seq_len_kv_cache); + } + 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 @@ -409,9 +625,13 @@ struct FMHAFwdMainloop< } page_idx = next_page_idx; - next_page_idx = K + 1; + int next_k = K + 1; + bool has_next_k = next_k < blk_k1 && next_k < kblocks_cache; + next_page_idx = next_k; if constexpr (PagedKV) { - next_page_idx = get_physical_k_tile(next_page_idx, l_coord, seq_len_kv_cache); + if (has_next_k) { + next_page_idx = get_physical_k_tile(next_k, l_coord, seq_len_kv_cache); + } } /* GEMM 1: S = K * Q */ @@ -434,25 +654,16 @@ struct FMHAFwdMainloop< /* 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; + // Need to get global col and row indices to mask the elements + 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 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); - } + for (int i = 0; i < tSrS.size(); ++i) { + int row_idx = get<0>(cS_thread(i)); + int col_idx = get<1>(cS_thread(i)); + if (row_idx < col_idx - full_tile_offset) { + tSrS(i) = ElementS(-INFINITY); } } } @@ -513,8 +724,10 @@ struct FMHAFwdMainloop< } /* K prefetch */ - for (int D = 0; D < size<4>(pKgK); D++) { - prefetch(prefetch_k_cache, pKgK_cache(_, _, _, next_page_idx, D)); + if (has_next_k) { + for (int D = 0; D < size<4>(pKgK); D++) { + prefetch(prefetch_k_cache, pKgK_cache(_, _, _, next_page_idx, D)); + } } barrier_wait(ScopeWorkgroup); 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..1c9202c76 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 @@ -242,6 +242,25 @@ class XeFMHAFwdKernel { } } + CUTLASS_DEVICE + int get_k_new_len(MainloopParams const& mainloop, int batch) { + if constexpr (CollectiveMainloop::AppendKV) { + if (mainloop.append.ptr_K_new == nullptr || mainloop.append.ptr_V_new == nullptr || + mainloop.append.ptr_cache_seqlens == nullptr || mainloop.append.total_k_new <= 0 || + (mainloop.append.ptr_cu_seqlens_k_new == nullptr && mainloop.append.seq_len_kv_new <= 0)) { + return 0; + } + if (mainloop.append.ptr_cu_seqlens_k_new != nullptr) { + return mainloop.append.ptr_cu_seqlens_k_new[batch + 1] - mainloop.append.ptr_cu_seqlens_k_new[batch]; + } + return mainloop.append.seq_len_kv_new; + } else { + (void)mainloop; + (void)batch; + return 0; + } + } + CUTLASS_DEVICE void operator()(Params const& params, char* smem_buf) { using namespace sycl::ext::oneapi::this_work_item; @@ -275,6 +294,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 = get_k_new_len(params.mainloop, idx_b); + if (seq_k_new > 0) { + seq_k_eff = params.mainloop.append.ptr_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 +312,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 +323,28 @@ 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; + if constexpr (CollectiveMainloop::AppendKV) { + int const seq_k_new = get_k_new_len(params.mainloop, idx_b); + if (seq_k_new > 0) { + append_store_len = seq_k_new; + 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.append.ptr_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 @@ -410,10 +452,13 @@ 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, 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..a8b9023e6 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 @@ -49,6 +49,79 @@ void FmhaPrefillRunner<@HEAD_DIM@>::operator()(const Arguments& params) const { using SubgroupLayoutQK = cute::Layout, cute::_1, cute::_1>, cute::Stride>; +#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_runner.hpp b/src/sycl/kernels/flash_attention_v2/xe_fmha_fwd_prefill_runner.hpp index 8789e8dc3..c6ab6cd65 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 @@ -110,6 +110,7 @@ struct Arguments { // If provided, the actual length of each q/k sequence. int* __restrict__ seqused_q; int* __restrict__ seqused_k; + int* __restrict__ cache_seqlens_old; // The stride between rows of Oaccum. int64_t oaccum_split_stride; @@ -293,6 +294,22 @@ 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.append.ptr_K_new = static_cast(params.knew_ptr); + mainloop_args.append.ptr_V_new = static_cast(params.vnew_ptr); + mainloop_args.append.ptr_cu_seqlens_k_new = params.cu_seqlens_knew; + mainloop_args.append.ptr_cache_seqlens = params.cache_seqlens_old; + mainloop_args.append.seq_len_kv_new = params.seqlen_knew; + mainloop_args.append.total_k_new = params.total_knew; + } + typename FMHAPrefillKernel::Arguments arguments{ { shape, @@ -313,14 +330,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 +391,7 @@ struct FMHAConfig { decltype(cutlass::fmha::collective::get_sg_layout_pv(SubgroupLayoutQK{})), SubgroupLayoutPV_>; - template + template 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 +441,9 @@ struct FMHAConfig { GmemTiledCopyV, GmemTiledCopyK_cache, GmemTiledCopyV_cache, - LocalMask>; + LocalMask, + false, + AppendKV>; // Epilogue using CollectiveEpilogue = @@ -460,14 +472,36 @@ struct FMHAConfig { // Paged KV cache: the page table encodes absolute KV positions. 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_k > 0, "paged prefill requires positive max sequence lengths"); + TORCH_CHECK(params.total_q > 0 && params.total_k > 0, "paged prefill requires positive total sequence lengths"); + bool const has_append = params.total_knew > 0 && params.knew_ptr != nullptr && params.vnew_ptr != nullptr && + params.cache_seqlens_old != nullptr; + // template + if (has_append) { + return run(params); + } + return run(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_k > 0, "non-paged prefill requires positive max sequence lengths"); + TORCH_CHECK(params.total_q > 0 && params.total_k > 0, "non-paged prefill requires positive total sequence lengths"); + bool const has_append = params.total_knew > 0 && params.knew_ptr != nullptr && params.vnew_ptr != nullptr && + params.cache_seqlens_old != nullptr; + // template + if (has_append) { + return run(params); + } + return run(params); } static int run(const Arguments& params) { diff --git a/src/torch_extension_sycl.cc b/src/torch_extension_sycl.cc index 30e594eb8..73c3d4d07 100644 --- a/src/torch_extension_sycl.cc +++ b/src/torch_extension_sycl.cc @@ -177,6 +177,40 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { " int sm_margin," " Tensor(a!)? out=None) -> (Tensor(a!), Tensor, Tensor, Tensor)"); m.impl("fwd", torch::kXPU, make_pytorch_shim(&mha_fwd)); + m.def( + "fwd_appendkv(Tensor q," + " Tensor k," + " Tensor v," + " Tensor? q_v," + " Tensor cu_seqlens_q," + " Tensor cu_seqlens_k," + " int max_seqlen_q," + " int max_seqlen_k," + " Tensor? page_table," + " Tensor? kv_batch_idx," + " Tensor? leftpad_k," + " Tensor? rotary_cos," + " Tensor? rotary_sin," + " Tensor? seqlens_rotary," + " Tensor? q_descale," + " Tensor? k_descale," + " Tensor? v_descale," + " float softmax_scale," + " Tensor? sinks," + " bool is_causal," + " int window_size_left," + " int window_size_right," + " float softcap," + " bool is_rotary_interleaved," + " Tensor? scheduler_metadata," + " int num_kv_splits," + " bool? pack_gqa," + " int sm_margin," + " 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_appendkv", torch::kXPU, make_pytorch_shim(&mha_fwd_appendkv)); #endif // USE_FMHA #ifdef USE_MLA 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, From 8b0f5cc62e179a459715163119a6bfc42245d36d Mon Sep 17 00:00:00 2001 From: Yuankun Shi Date: Fri, 17 Jul 2026 19:05:20 -0700 Subject: [PATCH 02/16] Tune hd64 FMHA prefill mixed batches --- .../xe_fmha_fwd_prefill_kernel.cpp.in | 75 +++++++++++++++++++ 1 file changed, 75 insertions(+) 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 a8b9023e6..31e757bb7 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 @@ -49,6 +49,81 @@ void FmhaPrefillRunner<@HEAD_DIM@>::operator()(const Arguments& params) const { using SubgroupLayoutQK = cute::Layout, cute::_1, cute::_1>, cute::Stride>; +#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<64>, cute::_32>; + using SparseTileShapePV = cute::Shape, cute::_32, cute::Int<64>>; + 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 <= 256 && + static_cast(params.total_q) * 4 <= static_cast(params.b) * params.seqlen_q && + params.d == 64 && params.dv == 64) { + using MixedTileShapeQK = cute::Shape, cute::Int<64>, cute::_32>; + using MixedTileShapePV = cute::Shape, cute::_32, cute::Int<64>>; + using MixedTileShapeOutput = cute::Shape, cute::Int<64>>; + using MixedSubgroupLayoutQK = 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 && From e311036f88134b58747d0d52992f2d25231511c3 Mon Sep 17 00:00:00 2001 From: Yuankun Shi Date: Fri, 17 Jul 2026 20:00:53 -0700 Subject: [PATCH 03/16] Tune hd64 sparse mixed KV tile --- .../flash_attention_v2/xe_fmha_fwd_prefill_kernel.cpp.in | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 31e757bb7..5f7f64865 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 @@ -53,8 +53,8 @@ void FmhaPrefillRunner<@HEAD_DIM@>::operator()(const Arguments& params) const { 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<64>, cute::_32>; - using SparseTileShapePV = cute::Shape, cute::_32, cute::Int<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>; From f107c1e06f5e09f34e3c285eedf6a1990554a5d6 Mon Sep 17 00:00:00 2001 From: Yuankun Shi Date: Fri, 17 Jul 2026 20:57:42 -0700 Subject: [PATCH 04/16] Drop unstable hd64 balanced mixed tile path --- .../xe_fmha_fwd_prefill_kernel.cpp.in | 19 ------------------- 1 file changed, 19 deletions(-) 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 5f7f64865..2cf0efffc 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 @@ -69,25 +69,6 @@ void FmhaPrefillRunner<@HEAD_DIM@>::operator()(const Arguments& params) const { return; } - if (params.seqlen_q <= 256 && - static_cast(params.total_q) * 4 <= static_cast(params.b) * params.seqlen_q && - params.d == 64 && params.dv == 64) { - using MixedTileShapeQK = cute::Shape, cute::Int<64>, cute::_32>; - using MixedTileShapePV = cute::Shape, cute::_32, cute::Int<64>>; - using MixedTileShapeOutput = cute::Shape, cute::Int<64>>; - using MixedSubgroupLayoutQK = 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>>; From aa07ff7b07994a91af508e20b4e4c7507b3f2f82 Mon Sep 17 00:00:00 2001 From: Yuankun Shi Date: Fri, 17 Jul 2026 21:45:21 -0700 Subject: [PATCH 05/16] Tune hd64 sparse mixed subgroup layout --- .../flash_attention_v2/xe_fmha_fwd_prefill_kernel.cpp.in | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 2cf0efffc..391cee2c6 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 @@ -57,7 +57,7 @@ void FmhaPrefillRunner<@HEAD_DIM@>::operator()(const Arguments& params) const { 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>; + 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, { From 574851a756d0efe47c3961eec92eb5af4155c8e8 Mon Sep 17 00:00:00 2001 From: Yuankun Shi Date: Sun, 26 Jul 2026 13:54:10 -0700 Subject: [PATCH 06/16] Tune hd256 FMHA prefill tile Port the q128/k64/sg16 paged tile configuration from cutlass-sycl commit 9bf71715. --- src/FMHAPrefillXe20.cmake | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/FMHAPrefillXe20.cmake b/src/FMHAPrefillXe20.cmake index bdd1c4920..e9ea9093a 100644 --- a/src/FMHAPrefillXe20.cmake +++ b/src/FMHAPrefillXe20.cmake @@ -46,9 +46,9 @@ set(FMHA_PREFILL_TILED_Q_192 256) set(FMHA_PREFILL_TILED_KV_192 64) set(FMHA_PREFILL_NUM_SG_192 32) -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) set(FMHA_PREFILL_TILED_Q_512 256) set(FMHA_PREFILL_TILED_KV_512 64) From d41ee92a055e316bf23a09110f6274606ef478a5 Mon Sep 17 00:00:00 2001 From: Yuankun Shi Date: Sun, 26 Jul 2026 14:08:27 -0700 Subject: [PATCH 07/16] Tune hd192 FMHA prefill tile --- src/FMHAPrefillXe20.cmake | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/FMHAPrefillXe20.cmake b/src/FMHAPrefillXe20.cmake index e9ea9093a..903a3d778 100644 --- a/src/FMHAPrefillXe20.cmake +++ b/src/FMHAPrefillXe20.cmake @@ -42,9 +42,9 @@ 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 128) set(FMHA_PREFILL_TILED_KV_256 64) From 4aa21ca9873ab15ca1f03c150840e0c269c52403 Mon Sep 17 00:00:00 2001 From: Yuankun Shi Date: Sun, 26 Jul 2026 14:58:41 -0700 Subject: [PATCH 08/16] Tune hd96 FMHA prefill tiles --- src/FMHAPrefillXe20.cmake | 34 +++++++++- .../xe_fmha_fwd_prefill_kernel.cpp.in | 28 ++++++++ .../xe_fmha_fwd_prefill_nopage_kernel.cpp.in | 65 +++++++++++++++---- 3 files changed, 112 insertions(+), 15 deletions(-) diff --git a/src/FMHAPrefillXe20.cmake b/src/FMHAPrefillXe20.cmake index 903a3d778..00f864868 100644 --- a/src/FMHAPrefillXe20.cmake +++ b/src/FMHAPrefillXe20.cmake @@ -31,7 +31,7 @@ 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 128) set(FMHA_PREFILL_TILED_KV_128 64) @@ -71,9 +71,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) @@ -131,6 +141,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/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 391cee2c6..bb1316b82 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,20 @@ void FmhaPrefillRunner<@HEAD_DIM@>::operator()(const Arguments& params) const { using SubgroupLayoutQK = cute::Layout, cute::_1, cute::_1>, cute::Stride>; +#if @HEAD_DIM@ == 96 + // Dense hkv8 batches already provide enough occupancy and run faster with the original sg8 tile. + 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@ == 64 if (params.seqlen_q <= 256 && static_cast(params.total_q) * 100 <= static_cast(params.b) * params.seqlen_q * 15 && 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..9b2fad876 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,54 @@ void FmhaPrefillNpRunner<@HEAD_DIM@>::operator()(const Arguments& params) const }); } +} // namespace + +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>; + +#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 From 49183d2df7dfd6b17a87f4aad5ae485c8df26c0a Mon Sep 17 00:00:00 2001 From: Yuankun Shi Date: Mon, 27 Jul 2026 10:17:11 -0700 Subject: [PATCH 09/16] Optimize q256 AppendKV prefill --- src/FMHAPrefillXe20.cmake | 26 ++++++++ .../collective/xe_fmha_fwd_mainloop.hpp | 65 +++++++++++++------ .../xe_fmha_fwd_prefill_kernel.cpp.in | 20 ++++++ 3 files changed, 91 insertions(+), 20 deletions(-) diff --git a/src/FMHAPrefillXe20.cmake b/src/FMHAPrefillXe20.cmake index 00f864868..07e20c87e 100644 --- a/src/FMHAPrefillXe20.cmake +++ b/src/FMHAPrefillXe20.cmake @@ -50,6 +50,19 @@ set(FMHA_PREFILL_TILED_Q_256 128) set(FMHA_PREFILL_TILED_KV_256 64) set(FMHA_PREFILL_NUM_SG_256 16) +# A q256 tile reduces duplicate fused KV-cache stores for uniform 256-token +# appends while retaining the tuned default tiles for shorter and ragged Q. +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) set(FMHA_PREFILL_NUM_SG_512 32) @@ -121,6 +134,19 @@ foreach(HEAD_DIM ${FMHA_PREFILL_PAGED_HEAD_DIMS}) set(HD128_PAGED_LARGE_TILE_KV ${FMHA_PREFILL_HD128_LARGE_TILE_KV}) set(HD128_PAGED_LARGE_NUM_SG ${FMHA_PREFILL_HD128_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) 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 713ec7956..4c7a66dc9 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 @@ -346,6 +346,9 @@ struct FMHAFwdMainloop< } int const head_size = head_size_qk; int const vecs_per_token = head_size / kVecElems; + bool const flatten_token_vectors = new_len >= kMinMultiSgTokens && head_size == 64; + 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) { @@ -377,10 +380,12 @@ struct FMHAFwdMainloop< ((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; - 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) { + 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.append.ptr_K_new + src); @@ -388,6 +393,19 @@ struct FMHAFwdMainloop< *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.append.ptr_K_new + src); + StoreVec const v_value = *reinterpret_cast(params.append.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; } @@ -395,22 +413,14 @@ struct FMHAFwdMainloop< } } - 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 dst_row = single_row_base + new_tok; - if constexpr (PagedKV) { - if (!single_dst_page) { - int const dst_tok = cache_len_old + new_tok; - 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; - } - } - - 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) { + 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.append.ptr_K_new + src); @@ -418,6 +428,21 @@ struct FMHAFwdMainloop< *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.append.ptr_K_new + src); + StoreVec const v_value = *reinterpret_cast(params.append.ptr_V_new + src); + *reinterpret_cast(&K_dst(dst_row, d)) = k_value; + *reinterpret_cast(&V_dst(d, dst_row)) = v_value; + } + } } return; } 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 bb1316b82..dc3dd60f2 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 @@ -63,6 +63,26 @@ 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_knew > 0 && params.seqlen_q == @APPENDKV_Q256_Q@ && + params.seqlen_knew == @APPENDKV_Q256_Q@ && + params.total_q == params.b * @APPENDKV_Q256_Q@ && + params.total_knew == 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>; + FMHAConfig::run_paged(params); + return; + } +#endif + #if @HEAD_DIM@ == 96 // Dense hkv8 batches already provide enough occupancy and run faster with the original sg8 tile. if (params.b >= 8 && params.h_k >= 8) { From 2d09c4af6f47d5adf182dd76ec249844a9787954 Mon Sep 17 00:00:00 2001 From: Yuankun Shi Date: Mon, 27 Jul 2026 12:50:03 -0700 Subject: [PATCH 10/16] Optimize q256 AppendKV cache access --- .../collective/xe_fmha_fwd_mainloop.hpp | 297 +++++++++++------- .../kernel/xe_fmha_fwd_kernel.hpp | 47 ++- .../xe_fmha_fwd_prefill_kernel.cpp.in | 7 +- .../xe_fmha_fwd_prefill_runner.hpp | 34 +- 4 files changed, 260 insertions(+), 125 deletions(-) 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 4c7a66dc9..7e5b3934f 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 @@ -94,7 +94,9 @@ template < // 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 AppendKV_ = false> + bool AppendKV_ = false, + bool DirectAppendKV_ = false, + bool WideAppendKV_ = false> struct FMHAFwdMainloop { static_assert(cutlass::detail::dependent_false, "Could not find a mainloop specialization."); }; @@ -121,7 +123,9 @@ template < class TiledCopyV_cache_, bool LocalMask_, bool PackGQA_, - bool AppendKV_> + bool AppendKV_, + bool DirectAppendKV_, + bool WideAppendKV_> struct FMHAFwdMainloop< XeDefault, CausalMask_, @@ -142,7 +146,9 @@ struct FMHAFwdMainloop< TiledCopyV_cache_, LocalMask_, PackGQA_, - AppendKV_> { + AppendKV_, + DirectAppendKV_, + WideAppendKV_> { // // Type Aliases // @@ -221,6 +227,11 @@ struct FMHAFwdMainloop< // scale) inside the mainloop after the block-2D load. static constexpr bool Fp8KV = is_any_of_v; static constexpr bool AppendKV = AppendKV_; + static constexpr bool DirectAppendKV = DirectAppendKV_; + static constexpr bool WideAppendKV = WideAppendKV_; + static_assert(!DirectAppendKV || (AppendKV && PagedKV), "Direct AppendKV requires paged AppendKV"); + static_assert( + !WideAppendKV || (AppendKV && PagedKV && !DirectAppendKV), "Wide AppendKV requires fused paged AppendKV"); using AppendKVStorage = AppendKVParams; @@ -306,11 +317,14 @@ struct FMHAFwdMainloop< int kv_head, int num_heads_kv, int thr_id, - int append_store_len = -1) const { + int append_store_len = -1, + int append_store_begin = 0) const { if constexpr (AppendKV) { int const new_len_total = get_k_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 ? new_len_total : (append_store_len < new_len_total ? append_store_len : new_len_total); + append_store_len < 0 ? available_len : (append_store_len < available_len ? append_store_len : available_len); if (new_len <= 0) { return; } @@ -319,9 +333,10 @@ struct FMHAFwdMainloop< 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.append.ptr_cu_seqlens_k_new != nullptr ? params.append.ptr_cu_seqlens_k_new[batch] - : batch * params.append.seq_len_kv_new; - int const cache_len_old = params.append.ptr_cache_seqlens[batch]; + int const new_begin = (params.append.ptr_cu_seqlens_k_new != nullptr ? params.append.ptr_cu_seqlens_k_new[batch] + : batch * params.append.seq_len_kv_new) + + store_begin; + int const cache_len_old = params.append.ptr_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; @@ -333,9 +348,10 @@ struct FMHAFwdMainloop< // 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) { - // NHD bf16 rows are 16B-aligned when head_size is divisible by 8. - constexpr int kVecElems = 8; - using StoreVec = cutlass::ulonglong2; + // The q256 HD64 specialization uses 32B vectors to reduce flattened + // store index overhead; other paths retain the 16B transaction. + constexpr int kVecElems = WideAppendKV ? 16 : 8; + using StoreVec = cute::conditional_t; if (head_size_qk == head_size_vo && (head_size_qk % kVecElems) == 0) { // Large appends are bandwidth-bound, so split tokens across SGs; // small appends keep the old single-SG path to avoid control overhead. @@ -345,8 +361,8 @@ struct FMHAFwdMainloop< return; } int const head_size = head_size_qk; - int const vecs_per_token = head_size / kVecElems; 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; @@ -481,6 +497,7 @@ struct FMHAFwdMainloop< (void)num_heads_kv; (void)thr_id; (void)append_store_len; + (void)append_store_begin; } } @@ -506,6 +523,7 @@ struct FMHAFwdMainloop< 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 @@ -599,27 +617,26 @@ struct FMHAFwdMainloop< // ------ if constexpr (AppendKV) { - store_kv_new(K_cache_2D, V_cache_2D, l_coord, kv_head, num_heads_kv, thr_id, append_store_len); - barrier(); + store_kv_new( + K_cache_2D, V_cache_2D, l_coord, kv_head, num_heads_kv, thr_id, append_store_len, append_store_begin); + if constexpr (DirectAppendKV) { + constexpr int kTileKV = get<1>(TileShapeQK{}); + int const cache_len_old = params.append.ptr_cache_seqlens[l_coord]; + int const new_begin = params.append.ptr_cu_seqlens_k_new != nullptr + ? params.append.ptr_cu_seqlens_k_new[l_coord] + : l_coord * params.append.seq_len_kv_new; + if ((cache_len_old % kTileKV) != 0 || (new_begin % kTileKV) != 0) { + barrier(); + } + } else { + barrier(); + } } - /* 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; + /* Initialization steps for first block: Q prefetch, O init */ for (int D = 0; D < size<3>(pQgQ); D++) { prefetch(prefetch_q, pQgQ(_, _, _, D)); } - bool has_prefetch_k = blk_k0 < blk_k1 && blk_k0 < kblocks_cache; - if (has_prefetch_k) { - if constexpr (PagedKV) { - next_page_idx = get_physical_k_tile(blk_k0, l_coord, seq_len_kv_cache); - } - 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). @@ -639,47 +656,103 @@ 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 (DirectAppendKV) { + int const cache_len_old = params.append.ptr_cache_seqlens[l_coord]; + int const new_begin = params.append.ptr_cu_seqlens_k_new != nullptr ? params.append.ptr_cu_seqlens_k_new[l_coord] + : l_coord * params.append.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; - int next_k = K + 1; - bool has_next_k = next_k < blk_k1 && next_k < kblocks_cache; - next_page_idx = next_k; - if constexpr (PagedKV) { - if (has_next_k) { - next_page_idx = get_physical_k_tile(next_k, 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) { - // Need to get global col and row indices to mask the elements + 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); + } + } + } + + clear(tSrS); + CUTLASS_PRAGMA_UNROLL + 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); + } + + CUTLASS_PRAGMA_UNROLL + 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)); + } + } + + if constexpr (CausalMask) { + if (need_causal) { + 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)); + if (row_idx < col_idx - full_tile_offset) { + tSrS(i) = ElementS(-INFINITY); + } + } + } + } + + 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); @@ -687,75 +760,71 @@ struct FMHAFwdMainloop< for (int i = 0; i < tSrS.size(); ++i) { int row_idx = get<0>(cS_thread(i)); int col_idx = get<1>(cS_thread(i)); - if (row_idx < col_idx - full_tile_offset) { + 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); } } } - } - /* 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); - 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); + 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 < tSrS.size(); i++) { + tSrS(i) = sycl::fmin(tSrS(i), broadcast<1>(k_rem_mask, tSrS, i)); } } - } - /* 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); - } + auto rescale = softmax(K == blk_k0, tSrS, tA_max, tA_sum, qk_scale); + reorder(tSrS, tArP); + CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < tSrS.size(); i++) { - tSrS(i) = sycl::fmin(tSrS(i), broadcast<1>(k_rem_mask, tSrS, i)); + 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)); } - } - - /* 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); - /* 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) { - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < tArA.size() / VTiles; i++) { - tArA(_, _, _, VV)(i) *= broadcast<0>(rescale, tArA, i); + 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)); + } } } - cute::gemm(mma_pv, tArP, tArV, tArA(_, _, _, VV)); - } - /* K prefetch */ - if (has_next_k) { - for (int D = 0; D < size<4>(pKgK); D++) { - prefetch(prefetch_k_cache, pKgK_cache(_, _, _, next_page_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 (DirectAppendKV) { + 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_fwd_kernel.hpp b/src/sycl/kernels/flash_attention_v2/kernel/xe_fmha_fwd_kernel.hpp index 1c9202c76..dd7934135 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 @@ -102,6 +102,7 @@ class XeFMHAFwdKernel { using FragA = typename CollectiveMainloop::FragA; using FragARow = typename CollectiveMainloop::FragARow; + static constexpr bool DirectAppendKV = CollectiveMainloop::DirectAppendKV; // Tile scheduler derived types using TileScheduler = TileScheduler_; @@ -329,10 +330,29 @@ class XeFMHAFwdKernel { 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 = get_k_new_len(params.mainloop, idx_b); if (seq_k_new > 0) { append_store_len = seq_k_new; + if constexpr (DirectAppendKV) { + int const cache_len_old = params.mainloop.append.ptr_cache_seqlens[idx_b]; + int const new_begin = params.mainloop.append.ptr_cu_seqlens_k_new != nullptr + ? params.mainloop.append.ptr_cu_seqlens_k_new[idx_b] + : idx_b * params.mainloop.append.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. @@ -395,15 +415,23 @@ 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; + if constexpr (DirectAppendKV) { + 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(DirectAppendKV ? p.K : p.K_cache); + auto dcV = const_cast(DirectAppendKV ? 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); @@ -411,13 +439,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; @@ -440,8 +474,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, @@ -459,6 +493,7 @@ class XeFMHAFwdKernel { 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 dc3dd60f2..064e9b03d 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 @@ -77,8 +77,13 @@ void FmhaPrefillRunner<@HEAD_DIM@>::operator()(const Arguments& params) const { using AppendSubgroupLayoutQK = cute::Layout< cute::Shape, cute::_1, cute::_1>, cute::Stride>; +#if @HEAD_DIM@ == 64 + FMHAConfig::run_paged_wide_append(params); +#else FMHAConfig::run_paged(params); + AppendSubgroupLayoutQK>::run_paged_direct_append(params); +#endif return; } #endif 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 c6ab6cd65..9bf906c39 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 @@ -315,9 +315,9 @@ struct PrefillRunner { shape, static_cast(params.q_ptr), stride_Q, - nullptr, + CollectiveMainloop::DirectAppendKV ? static_cast(params.knew_ptr) : nullptr, stride_K, - nullptr, + CollectiveMainloop::DirectAppendKV ? static_cast(params.vnew_ptr) : nullptr, stride_V, static_cast(params.o_ptr), stride_O, @@ -391,7 +391,14 @@ struct FMHAConfig { decltype(cutlass::fmha::collective::get_sg_layout_pv(SubgroupLayoutQK{})), SubgroupLayoutPV_>; - template + template < + bool isVarLen, + bool CachedKV, + bool PagedKV, + bool AppendKV, + class Scheduler, + bool DirectAppendKV = false, + bool WideAppendKV = false> 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. @@ -443,7 +450,9 @@ struct FMHAConfig { GmemTiledCopyV_cache, LocalMask, false, - AppendKV>; + AppendKV, + DirectAppendKV, + WideAppendKV>; // Epilogue using CollectiveEpilogue = @@ -488,6 +497,23 @@ struct FMHAConfig { return run(params); } + static int run_paged_direct_append(const Arguments& params) { + TORCH_CHECK(params.cu_seqlens_q != nullptr, "direct AppendKV requires cu_seqlens_q"); + TORCH_CHECK(params.cu_seqlens_k != nullptr, "direct AppendKV requires per-batch cache lengths"); + TORCH_CHECK(params.page_table != nullptr, "direct AppendKV requires page_table"); + TORCH_CHECK(params.page_size > 0, "direct AppendKV requires a positive page_size"); + TORCH_CHECK(params.max_num_pages_per_seq > 0, "direct AppendKV requires max_num_pages_per_seq"); + TORCH_CHECK( + params.total_knew > 0 && params.knew_ptr != nullptr && params.vnew_ptr != nullptr && + params.cache_seqlens_old != nullptr, + "direct AppendKV requires k_new, v_new, and old cache lengths"); + return run(params); + } + + static int run_paged_wide_append(const Arguments& params) { + return run(params); + } + // Non-paged (contiguous ragged) KV cache: addressed via cu_seqlens_k offsets. static int run_nopaged(const Arguments& params) { TORCH_CHECK(params.cu_seqlens_q != nullptr, "non-paged prefill requires cu_seqlens_q"); From b6c23bd5892f4f8d503fe4bbcd6db7f4018f1252 Mon Sep 17 00:00:00 2001 From: Yuankun Shi Date: Tue, 28 Jul 2026 21:51:19 -0700 Subject: [PATCH 11/16] Add FMHA prefill AppendKV wiring --- include/sgl_flash_kernel_ops.h | 31 --- python/sgl_kernel/flash_attn.py | 101 +++------ src/FMHAPrefillXe20.cmake | 2 - src/sycl/flash_attention.cpp | 205 ++++-------------- .../collective/xe_fmha_fwd_mainloop.hpp | 84 ++++--- .../kernel/xe_fmha_fwd_kernel.hpp | 24 +- .../xe_fmha_fwd_prefill_kernel.cpp.in | 1 - .../xe_fmha_fwd_prefill_nopage_kernel.cpp.in | 2 - .../xe_fmha_fwd_prefill_runner.hpp | 14 +- src/torch_extension_sycl.cc | 33 +-- 10 files changed, 140 insertions(+), 357 deletions(-) diff --git a/include/sgl_flash_kernel_ops.h b/include/sgl_flash_kernel_ops.h index 16842c400..c3a28c445 100644 --- a/include/sgl_flash_kernel_ops.h +++ b/include/sgl_flash_kernel_ops.h @@ -74,37 +74,6 @@ std::tuple mha_fwd( int num_kv_splits, std::optional pack_gqa_, int const sm_margin, - std::optional& out_); - -std::tuple mha_fwd_appendkv( - const at::Tensor& q, - const at::Tensor& k, - const at::Tensor& v, - std::optional& q_v_, - const at::Tensor& cu_seqlens_q, - const at::Tensor& cu_seqlens_k, - int max_seqlen_q, - int max_seqlen_k, - std::optional& page_table, - std::optional& kv_batch_idx_, - std::optional& leftpad_k_, - std::optional& rotary_cos_, - std::optional& rotary_sin_, - std::optional& seqlens_rotary_, - std::optional& q_descale_, - std::optional& k_descale_, - std::optional& v_descale_, - float const softmax_scale, - std::optional& sinks, - bool is_causal, - int window_size_left, - int window_size_right, - float const softcap, - bool const is_rotary_interleaved, - std::optional& scheduler_metadata_, - int num_kv_splits, - std::optional pack_gqa_, - int const sm_margin, std::optional& out_, std::optional& k_new_, std::optional& v_new_, diff --git a/python/sgl_kernel/flash_attn.py b/python/sgl_kernel/flash_attn.py index bf2607af9..930d4be32 100644 --- a/python/sgl_kernel/flash_attn.py +++ b/python/sgl_kernel/flash_attn.py @@ -265,78 +265,47 @@ def flash_attn_with_kvcache( 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_appendkv.default( - q, - k_cache, - v_cache, - qv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - native_max_seqlen_k, - page_table, - cache_batch_idx, - cache_leftpad, - rotary_cos, - rotary_sin, - rotary_seqlens, - q_descale, - k_descale, - v_descale, - softmax_scale, - sinks, - causal, - window_size[0], - window_size[1], - softcap, - rotary_interleaved, - scheduler_metadata, - num_splits, - pack_gqa, - sm_margin, - out, - k, - v, - cu_seqlens_k_new, - ) - else: - out, softmax_lse, *rest = torch.ops.sgl_kernel.fwd.default( - q, - k_cache, - v_cache, - qv, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - 1, - page_table, - cache_batch_idx, - cache_leftpad, - rotary_cos, - rotary_sin, - rotary_seqlens, - q_descale, - k_descale, - v_descale, - softmax_scale, - sinks, - causal, - window_size[0], - window_size[1], - softcap, - rotary_interleaved, - scheduler_metadata, - num_splits, - pack_gqa, - sm_margin, - out, - ) + out, softmax_lse, *rest = torch.ops.sgl_kernel.fwd.default( + q, + k_cache, + v_cache, + qv, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + native_max_seqlen_k, + page_table, + cache_batch_idx, + cache_leftpad, + rotary_cos, + rotary_sin, + rotary_seqlens, + q_descale, + k_descale, + v_descale, + softmax_scale, + sinks, + causal, + window_size[0], + window_size[1], + softcap, + rotary_interleaved, + scheduler_metadata, + num_splits, + 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 07e20c87e..b5f76f209 100644 --- a/src/FMHAPrefillXe20.cmake +++ b/src/FMHAPrefillXe20.cmake @@ -50,8 +50,6 @@ set(FMHA_PREFILL_TILED_Q_256 128) set(FMHA_PREFILL_TILED_KV_256 64) set(FMHA_PREFILL_NUM_SG_256 16) -# A q256 tile reduces duplicate fused KV-cache stores for uniform 256-token -# appends while retaining the tuned default tiles for shorter and ragged Q. 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) diff --git a/src/sycl/flash_attention.cpp b/src/sycl/flash_attention.cpp index 155530f09..23d338375 100644 --- a/src/sycl/flash_attention.cpp +++ b/src/sycl/flash_attention.cpp @@ -870,7 +870,7 @@ std::vector mha_fwd_nopage( return {out, softmax_lse, out_accum, softmax_lse_accum}; } -std::vector mha_fwd_appendkv( +std::vector mha_fwd( const at::Tensor& q, // (b, s_q, h, d) or (total_q, h, d) if there is cu_seqlens_q const at::Tensor& k, // (b_k, s_k, h_k, d) or (total_k, h_k, d) if there is cu_seqlens_k or (num_pages, page_size, // h_k, d) if there is page_table. @@ -1259,75 +1259,6 @@ std::vector mha_fwd_appendkv( return {out, softmax_lse, out_accum, softmax_lse_accum}; } -std::vector mha_fwd( - const at::Tensor& q, // (b, s_q, h, d) or (total_q, h, d) if there is cu_seqlens_q - const at::Tensor& k, // (b_k, s_k, h_k, d) or (total_k, h_k, d) if there is cu_seqlens_k or (num_pages, page_size, - // h_k, d) if there is page_table. - const at::Tensor& v, // (b_k, s_k, h_k, dv) or (total_k, h_k, dv) if there is cu_seqlens_k or (num_pages, - // page_size, h_k, dv) if there is page_table. - std::optional& q_v_, // (b, s_q, h, dv) or (total_q_new, h, dv) if there is cu_seqlens_q - const at::Tensor& cu_seqlens_q, // b+1 - const at::Tensor& cu_seqlens_k, // b+1 - int max_seqlen_q, - int max_seqlen_k, - std::optional& page_table, // (b_k, max_num_pages_per_seq) - std::optional& kv_batch_idx_, // b. indices to index into the KV cache - std::optional& leftpad_k_, // b - std::optional& rotary_cos_, // seqlen_ro x (rotary_dim / 2) - std::optional& rotary_sin_, // seqlen_ro x (rotary_dim / 2) - std::optional& seqlens_rotary_, // b - std::optional& q_descale_, // (b, h_k), not (b, h) - std::optional& k_descale_, // (b, h_k) - std::optional& v_descale_, // (b, h_k) - const float softmax_scale_, - std::optional& sinks_, - bool is_causal, - int window_size_left, - int window_size_right, - float const softcap, - bool const is_rotary_interleaved, // if true, rotary combines indices 0 & 1, else indices 0 & rotary_dim / 2 - std::optional& scheduler_metadata_, // (b + 1) - int num_splits, - std::optional pack_gqa_, - int const sm_margin, - std::optional out_opt = std::nullopt, - std::optional skip_batch_mask_opt = std::nullopt) { - return mha_fwd_appendkv( - 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_splits, - pack_gqa_, - sm_margin, - std::move(out_opt), - std::move(skip_batch_mask_opt), - std::nullopt, - std::nullopt, - std::nullopt); -} - } // namespace prefill namespace chunkprefill { @@ -1370,7 +1301,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 @@ -1381,12 +1315,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."); @@ -1434,7 +1370,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. @@ -1473,7 +1442,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. @@ -1541,92 +1513,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 { - // Prefill path + // Chunk prefill path // Paged attn with max_seqlen_q > 1 and batch_size > 1 - return dispatch(prefill::mha_fwd, std::nullopt); - } -} - -std::tuple mha_fwd_appendkv( - const at::Tensor& q, // (total_q, h, d) — ragged 3D - const at::Tensor& k, // (total_k, h_k, d) if non-paged, or (num_pages, page_size, h_k, d) if paged - const at::Tensor& v, // (total_k, h_k, dv) if non-paged, or (num_pages, page_size, h_k, dv) if paged - std::optional& q_v_, // (total_q, h, dv) — not yet supported - const at::Tensor& cu_seqlens_q, // b+1 - const at::Tensor& cu_seqlens_k, // b+1 - int max_seqlen_q, - int max_seqlen_k, - std::optional& page_table, // (b_k, max_num_pages_per_seq) - std::optional& kv_batch_idx_, // b. indices to index into the KV cache - std::optional& leftpad_k_, // b - std::optional& rotary_cos_, // seqlen_ro x (rotary_dim / 2) - std::optional& rotary_sin_, // seqlen_ro x (rotary_dim / 2) - std::optional& seqlens_rotary_, // b - std::optional& q_descale_, // (b, h_k), not (b, h) - std::optional& k_descale_, // (b, h_k) - std::optional& v_descale_, // (b, h_k) - const float softmax_scale_, - std::optional& sinks_, - bool is_causal, - int window_size_left, - int window_size_right, - float const softcap, - bool const is_rotary_interleaved, // if true, rotary combines indices 0 & 1, else indices 0 & rotary_dim / 2 - std::optional& scheduler_metadata_, // (b + 1) - int num_kv_splits, - std::optional pack_gqa_, - int const sm_margin, - 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)"); - if (out_.has_value()) { - const at::Tensor& out_val = out_.value(); - TORCH_CHECK(out_val.scalar_type() == q.scalar_type(), "out dtype must match q dtype"); - TORCH_CHECK( - out_val.dim() == 3 && out_val.size(0) == q.size(0) && out_val.size(1) == q.size(1) && - out_val.size(2) == v.size(-1), - "out shape must be [total_q, num_heads, head_size_v]"); - TORCH_CHECK(out_val.device() == q.device(), "out must be on the same device as q"); - TORCH_CHECK(out_val.stride(-1) == 1, "out must have a contiguous last dimension"); + return dispatch(chunkprefill::mha_fwd, k_new_, v_new_, cu_seqlens_k_new_); } - auto to_tuple = [](std::vector v) { return std::make_tuple(v[0], v[1], v[2], v[3]); }; - return to_tuple(prefill::mha_fwd_appendkv( - 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_, - std::nullopt, - 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 7e5b3934f..b0a0583eb 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 @@ -55,21 +55,6 @@ using namespace cute; ///////////////////////////////////////////////////////////////////////////////////////////////// -template -struct AppendKVParams {}; - -template -struct AppendKVParams { - ElementK_ const* ptr_K_new = nullptr; - ElementV_ const* ptr_V_new = nullptr; - int const* ptr_cu_seqlens_k_new = nullptr; - int const* ptr_cache_seqlens = nullptr; - int seq_len_kv_new = 0; - int total_k_new = 0; -}; - -///////////////////////////////////////////////////////////////////////////////////////////////// - template < class DispatchPolicy_, bool CausalMask_, @@ -232,8 +217,6 @@ struct FMHAFwdMainloop< static_assert(!DirectAppendKV || (AppendKV && PagedKV), "Direct AppendKV requires paged AppendKV"); static_assert( !WideAppendKV || (AppendKV && PagedKV && !DirectAppendKV), "Wide AppendKV requires fused paged AppendKV"); - using AppendKVStorage = - AppendKVParams; // User-facing arguments struct Arguments { @@ -243,7 +226,12 @@ struct FMHAFwdMainloop< int max_num_pages_per_seq = 0; int window_size_left = -1; int window_size_right = -1; - AppendKVStorage append{}; + 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_cache_seqlens = nullptr; + int seq_len_kv_new = 0; + int total_k_new = 0; }; // Kernel-facing parameters @@ -270,7 +258,12 @@ struct FMHAFwdMainloop< args.max_num_pages_per_seq, args.window_size_left, args.window_size_right, - args.append}; + args.ptr_K_new, + args.ptr_V_new, + args.ptr_cu_seqlens_k_new, + args.ptr_cache_seqlens, + args.seq_len_kv_new, + args.total_k_new}; } CUTLASS_HOST_DEVICE static bool can_implement(Arguments const&) { @@ -294,15 +287,15 @@ struct FMHAFwdMainloop< CUTLASS_DEVICE int get_k_new_len(int batch) const { if constexpr (AppendKV) { - if (params.append.ptr_K_new == nullptr || params.append.ptr_V_new == nullptr || - params.append.ptr_cache_seqlens == nullptr || params.append.total_k_new <= 0 || - (params.append.ptr_cu_seqlens_k_new == nullptr && params.append.seq_len_kv_new <= 0)) { + if (params.ptr_K_new == nullptr || params.ptr_V_new == nullptr || + params.ptr_cache_seqlens == nullptr || params.total_k_new <= 0 || + (params.ptr_cu_seqlens_k_new == nullptr && params.seq_len_kv_new <= 0)) { return 0; } - if (params.append.ptr_cu_seqlens_k_new != nullptr) { - return params.append.ptr_cu_seqlens_k_new[batch + 1] - params.append.ptr_cu_seqlens_k_new[batch]; + if (params.ptr_cu_seqlens_k_new != nullptr) { + return params.ptr_cu_seqlens_k_new[batch + 1] - params.ptr_cu_seqlens_k_new[batch]; } - return params.append.seq_len_kv_new; + return params.seq_len_kv_new; } else { (void)batch; return 0; @@ -333,10 +326,10 @@ struct FMHAFwdMainloop< 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.append.ptr_cu_seqlens_k_new != nullptr ? params.append.ptr_cu_seqlens_k_new[batch] - : batch * params.append.seq_len_kv_new) + + 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.append.ptr_cache_seqlens[batch] + store_begin; + int const cache_len_old = params.ptr_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; @@ -404,8 +397,8 @@ struct FMHAFwdMainloop< 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.append.ptr_K_new + src); - StoreVec const v_value = *reinterpret_cast(params.append.ptr_V_new + src); + 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; } @@ -416,8 +409,8 @@ struct FMHAFwdMainloop< 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.append.ptr_K_new + src); - StoreVec const v_value = *reinterpret_cast(params.append.ptr_V_new + src); + 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; } @@ -439,8 +432,8 @@ struct FMHAFwdMainloop< ((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.append.ptr_K_new + src); - StoreVec const v_value = *reinterpret_cast(params.append.ptr_V_new + src); + 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; } @@ -453,8 +446,8 @@ struct FMHAFwdMainloop< 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.append.ptr_K_new + src); - StoreVec const v_value = *reinterpret_cast(params.append.ptr_V_new + src); + 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; } @@ -481,12 +474,12 @@ struct FMHAFwdMainloop< 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.append.ptr_K_new[src]; + 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.append.ptr_V_new[src]; + V_dst(d, dst_row) = params.ptr_V_new[src]; } } } else { @@ -621,10 +614,10 @@ struct FMHAFwdMainloop< K_cache_2D, V_cache_2D, l_coord, kv_head, num_heads_kv, thr_id, append_store_len, append_store_begin); if constexpr (DirectAppendKV) { constexpr int kTileKV = get<1>(TileShapeQK{}); - int const cache_len_old = params.append.ptr_cache_seqlens[l_coord]; - int const new_begin = params.append.ptr_cu_seqlens_k_new != nullptr - ? params.append.ptr_cu_seqlens_k_new[l_coord] - : l_coord * params.append.seq_len_kv_new; + int const cache_len_old = params.ptr_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; if ((cache_len_old % kTileKV) != 0 || (new_begin % kTileKV) != 0) { barrier(); } @@ -633,7 +626,6 @@ struct FMHAFwdMainloop< } } - /* Initialization steps for first block: Q prefetch, O init */ for (int D = 0; D < size<3>(pQgQ); D++) { prefetch(prefetch_q, pQgQ(_, _, _, D)); } @@ -662,9 +654,9 @@ struct FMHAFwdMainloop< int direct_block0 = kblocks_total; int direct_source_block0 = 0; if constexpr (DirectAppendKV) { - int const cache_len_old = params.append.ptr_cache_seqlens[l_coord]; - int const new_begin = params.append.ptr_cu_seqlens_k_new != nullptr ? params.append.ptr_cu_seqlens_k_new[l_coord] - : l_coord * params.append.seq_len_kv_new; + int const cache_len_old = params.ptr_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; 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 dd7934135..82c009d7d 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 @@ -246,15 +246,15 @@ class XeFMHAFwdKernel { CUTLASS_DEVICE int get_k_new_len(MainloopParams const& mainloop, int batch) { if constexpr (CollectiveMainloop::AppendKV) { - if (mainloop.append.ptr_K_new == nullptr || mainloop.append.ptr_V_new == nullptr || - mainloop.append.ptr_cache_seqlens == nullptr || mainloop.append.total_k_new <= 0 || - (mainloop.append.ptr_cu_seqlens_k_new == nullptr && mainloop.append.seq_len_kv_new <= 0)) { + if (mainloop.ptr_K_new == nullptr || mainloop.ptr_V_new == nullptr || + mainloop.ptr_cache_seqlens == nullptr || mainloop.total_k_new <= 0 || + (mainloop.ptr_cu_seqlens_k_new == nullptr && mainloop.seq_len_kv_new <= 0)) { return 0; } - if (mainloop.append.ptr_cu_seqlens_k_new != nullptr) { - return mainloop.append.ptr_cu_seqlens_k_new[batch + 1] - mainloop.append.ptr_cu_seqlens_k_new[batch]; + if (mainloop.ptr_cu_seqlens_k_new != nullptr) { + return mainloop.ptr_cu_seqlens_k_new[batch + 1] - mainloop.ptr_cu_seqlens_k_new[batch]; } - return mainloop.append.seq_len_kv_new; + return mainloop.seq_len_kv_new; } else { (void)mainloop; (void)batch; @@ -299,7 +299,7 @@ class XeFMHAFwdKernel { if constexpr (CollectiveMainloop::AppendKV) { int const seq_k_new = get_k_new_len(params.mainloop, idx_b); if (seq_k_new > 0) { - seq_k_eff = params.mainloop.append.ptr_cache_seqlens[idx_b] + seq_k_new; + seq_k_eff = params.mainloop.ptr_cache_seqlens[idx_b] + seq_k_new; } } // M extent of the Q/O tile: the packed GQA group for decode, otherwise the @@ -336,10 +336,10 @@ class XeFMHAFwdKernel { if (seq_k_new > 0) { append_store_len = seq_k_new; if constexpr (DirectAppendKV) { - int const cache_len_old = params.mainloop.append.ptr_cache_seqlens[idx_b]; - int const new_begin = params.mainloop.append.ptr_cu_seqlens_k_new != nullptr - ? params.mainloop.append.ptr_cu_seqlens_k_new[idx_b] - : idx_b * params.mainloop.append.seq_len_kv_new; + int const cache_len_old = params.mainloop.ptr_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) { @@ -356,7 +356,7 @@ class XeFMHAFwdKernel { 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.append.ptr_cache_seqlens[idx_b]; + int const cache_len_old = params.mainloop.ptr_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); 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 064e9b03d..015aab15d 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 @@ -89,7 +89,6 @@ void FmhaPrefillRunner<@HEAD_DIM@>::operator()(const Arguments& params) const { #endif #if @HEAD_DIM@ == 96 - // Dense hkv8 batches already provide enough occupancy and run faster with the original sg8 tile. 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>>; 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 9b2fad876..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 @@ -72,8 +72,6 @@ void run_nopaged_tile(const Arguments& params) { 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)"); 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 9bf906c39..78a014d70 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 @@ -302,12 +302,12 @@ struct PrefillRunner { params.window_size_left, params.window_size_right}; if constexpr (CollectiveMainloop::AppendKV) { - mainloop_args.append.ptr_K_new = static_cast(params.knew_ptr); - mainloop_args.append.ptr_V_new = static_cast(params.vnew_ptr); - mainloop_args.append.ptr_cu_seqlens_k_new = params.cu_seqlens_knew; - mainloop_args.append.ptr_cache_seqlens = params.cache_seqlens_old; - mainloop_args.append.seq_len_kv_new = params.seqlen_knew; - mainloop_args.append.total_k_new = params.total_knew; + mainloop_args.ptr_K_new = static_cast(params.knew_ptr); + mainloop_args.ptr_V_new = static_cast(params.vnew_ptr); + mainloop_args.ptr_cu_seqlens_k_new = params.cu_seqlens_knew; + mainloop_args.ptr_cache_seqlens = params.cache_seqlens_old; + mainloop_args.seq_len_kv_new = params.seqlen_knew; + mainloop_args.total_k_new = params.total_knew; } typename FMHAPrefillKernel::Arguments arguments{ @@ -490,7 +490,6 @@ struct FMHAConfig { TORCH_CHECK(params.total_q > 0 && params.total_k > 0, "paged prefill requires positive total sequence lengths"); bool const has_append = params.total_knew > 0 && params.knew_ptr != nullptr && params.vnew_ptr != nullptr && params.cache_seqlens_old != nullptr; - // template if (has_append) { return run(params); } @@ -523,7 +522,6 @@ struct FMHAConfig { TORCH_CHECK(params.total_q > 0 && params.total_k > 0, "non-paged prefill requires positive total sequence lengths"); bool const has_append = params.total_knew > 0 && params.knew_ptr != nullptr && params.vnew_ptr != nullptr && params.cache_seqlens_old != nullptr; - // template if (has_append) { return run(params); } diff --git a/src/torch_extension_sycl.cc b/src/torch_extension_sycl.cc index 73c3d4d07..ba0668ef4 100644 --- a/src/torch_extension_sycl.cc +++ b/src/torch_extension_sycl.cc @@ -175,42 +175,11 @@ 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)"); - m.impl("fwd", torch::kXPU, make_pytorch_shim(&mha_fwd)); - m.def( - "fwd_appendkv(Tensor q," - " Tensor k," - " Tensor v," - " Tensor? q_v," - " Tensor cu_seqlens_q," - " Tensor cu_seqlens_k," - " int max_seqlen_q," - " int max_seqlen_k," - " Tensor? page_table," - " Tensor? kv_batch_idx," - " Tensor? leftpad_k," - " Tensor? rotary_cos," - " Tensor? rotary_sin," - " Tensor? seqlens_rotary," - " Tensor? q_descale," - " Tensor? k_descale," - " Tensor? v_descale," - " float softmax_scale," - " Tensor? sinks," - " bool is_causal," - " int window_size_left," - " int window_size_right," - " float softcap," - " bool is_rotary_interleaved," - " Tensor? scheduler_metadata," - " int num_kv_splits," - " bool? pack_gqa," - " int sm_margin," " 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_appendkv", torch::kXPU, make_pytorch_shim(&mha_fwd_appendkv)); + m.impl("fwd", torch::kXPU, make_pytorch_shim(&mha_fwd)); #endif // USE_FMHA #ifdef USE_MLA From e23eec7c0ea65756fc86627b2f6ee4587ac4f788 Mon Sep 17 00:00:00 2001 From: Yuankun Shi Date: Wed, 29 Jul 2026 15:14:34 -0700 Subject: [PATCH 12/16] Fix FMHA prefill AppendKV wiring --- src/sycl/flash_attention.cpp | 63 +++++++------- .../collective/xe_fmha_fwd_mainloop.hpp | 71 ++++++++++------ .../kernel/xe_fmha_fwd_kernel.hpp | 29 ++----- .../xe_fmha_fwd_prefill_kernel.cpp.in | 6 +- .../xe_fmha_fwd_prefill_runner.hpp | 85 +++++++++++-------- 5 files changed, 133 insertions(+), 121 deletions(-) diff --git a/src/sycl/flash_attention.cpp b/src/sycl/flash_attention.cpp index 23d338375..6a13fa6f6 100644 --- a/src/sycl/flash_attention.cpp +++ b/src/sycl/flash_attention.cpp @@ -183,8 +183,7 @@ 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; @@ -781,12 +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.knew_ptr = nullptr; - params.vnew_ptr = nullptr; - params.cu_seqlens_knew = nullptr; - params.cache_seqlens_old = nullptr; - params.seqlen_knew = 0; - params.total_knew = 0; + 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(); @@ -795,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; @@ -820,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; @@ -1068,7 +1067,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; @@ -1113,7 +1112,7 @@ 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.total_kvcache = total_k; params.b_k = batch_size_k; params.dv = head_size_v; params.page_table = page_table.value().data_ptr(); @@ -1122,12 +1121,12 @@ std::vector mha_fwd( params.page_size = page_size; params.num_pages = num_pages; - params.knew_ptr = nullptr; - params.vnew_ptr = nullptr; - params.cu_seqlens_knew = nullptr; - params.cache_seqlens_old = nullptr; - params.seqlen_knew = 0; - params.total_knew = 0; + 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(); @@ -1138,20 +1137,20 @@ std::vector mha_fwd( 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_knew = 0; - int seqlen_knew = max_seqlen_k > 0 ? max_seqlen_k : max_seqlen_q; + int total_kvnew = 0; + int seqlen_kvnew = max_seqlen_k > 0 ? max_seqlen_k : max_seqlen_q; if (k_new.dim() == 3) { - total_knew = k_new.size(0); - CHECK_SHAPE(k_new, total_knew, num_heads_k, head_size); - CHECK_SHAPE(v_new, total_knew, num_heads_k, head_size_v); + 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_knew > 0, + 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_knew = batch_size * k_new_seqlen; - seqlen_knew = max_seqlen_k > 0 ? max_seqlen_k : k_new_seqlen; + 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); } @@ -1160,13 +1159,13 @@ std::vector mha_fwd( 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_knew = cu_seqlens_k_new.data_ptr(); + params.cu_seqlens_kvnew = cu_seqlens_k_new.data_ptr(); } - params.knew_ptr = k_new.data_ptr(); - params.vnew_ptr = v_new.data_ptr(); - params.cache_seqlens_old = cu_seqlens_k.data_ptr(); - params.seqlen_knew = seqlen_knew; - params.total_knew = total_knew; + 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()) { 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 b0a0583eb..d952dd013 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 @@ -229,7 +229,7 @@ struct FMHAFwdMainloop< 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_cache_seqlens = nullptr; + int const* ptr_kv_cache_seqlens = nullptr; int seq_len_kv_new = 0; int total_k_new = 0; }; @@ -261,7 +261,7 @@ struct FMHAFwdMainloop< args.ptr_K_new, args.ptr_V_new, args.ptr_cu_seqlens_k_new, - args.ptr_cache_seqlens, + args.ptr_kv_cache_seqlens, args.seq_len_kv_new, args.total_k_new}; } @@ -285,23 +285,35 @@ struct FMHAFwdMainloop< } CUTLASS_DEVICE - int get_k_new_len(int batch) const { + 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_cache_seqlens == nullptr || params.total_k_new <= 0 || - (params.ptr_cu_seqlens_k_new == nullptr && params.seq_len_kv_new <= 0)) { + if (params.ptr_K_new == nullptr || params.ptr_V_new == nullptr || params.ptr_kv_cache_seqlens == nullptr) { return 0; } - if (params.ptr_cu_seqlens_k_new != nullptr) { - return params.ptr_cu_seqlens_k_new[batch + 1] - params.ptr_cu_seqlens_k_new[batch]; - } - return params.seq_len_kv_new; + 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); + } + CUTLASS_DEVICE void store_kv_new( TensorK_cache2D const& K_cache_2D, @@ -313,7 +325,7 @@ struct FMHAFwdMainloop< int append_store_len = -1, int append_store_begin = 0) const { if constexpr (AppendKV) { - int const new_len_total = get_k_new_len(batch); + 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 = @@ -329,7 +341,7 @@ struct FMHAFwdMainloop< 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_cache_seqlens[batch] + 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; @@ -614,10 +626,9 @@ struct FMHAFwdMainloop< K_cache_2D, V_cache_2D, l_coord, kv_head, num_heads_kv, thr_id, append_store_len, append_store_begin); if constexpr (DirectAppendKV) { constexpr int kTileKV = get<1>(TileShapeQK{}); - int const cache_len_old = params.ptr_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; + 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; if ((cache_len_old % kTileKV) != 0 || (new_begin % kTileKV) != 0) { barrier(); } @@ -654,7 +665,7 @@ struct FMHAFwdMainloop< int direct_block0 = kblocks_total; int direct_source_block0 = 0; if constexpr (DirectAppendKV) { - int const cache_len_old = params.ptr_cache_seqlens[l_coord]; + 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; @@ -730,15 +741,25 @@ struct FMHAFwdMainloop< if constexpr (CausalMask) { if (need_causal) { - 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); + /* 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 i = 0; i < tSrS.size(); ++i) { - int row_idx = get<0>(cS_thread(i)); - int col_idx = get<1>(cS_thread(i)); - if (row_idx < col_idx - full_tile_offset) { - tSrS(i) = ElementS(-INFINITY); + 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); + } } } } 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 82c009d7d..7df3fdcb8 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 @@ -243,25 +243,6 @@ class XeFMHAFwdKernel { } } - CUTLASS_DEVICE - int get_k_new_len(MainloopParams const& mainloop, int batch) { - if constexpr (CollectiveMainloop::AppendKV) { - if (mainloop.ptr_K_new == nullptr || mainloop.ptr_V_new == nullptr || - mainloop.ptr_cache_seqlens == nullptr || mainloop.total_k_new <= 0 || - (mainloop.ptr_cu_seqlens_k_new == nullptr && mainloop.seq_len_kv_new <= 0)) { - return 0; - } - if (mainloop.ptr_cu_seqlens_k_new != nullptr) { - return mainloop.ptr_cu_seqlens_k_new[batch + 1] - mainloop.ptr_cu_seqlens_k_new[batch]; - } - return mainloop.seq_len_kv_new; - } else { - (void)mainloop; - (void)batch; - return 0; - } - } - CUTLASS_DEVICE void operator()(Params const& params, char* smem_buf) { using namespace sycl::ext::oneapi::this_work_item; @@ -297,9 +278,9 @@ class XeFMHAFwdKernel { 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 = get_k_new_len(params.mainloop, idx_b); + 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_cache_seqlens[idx_b] + seq_k_new; + 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 @@ -332,11 +313,11 @@ class XeFMHAFwdKernel { int append_store_len = -1; int append_store_begin = 0; if constexpr (CollectiveMainloop::AppendKV) { - int const seq_k_new = get_k_new_len(params.mainloop, idx_b); + 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 (DirectAppendKV) { - int const cache_len_old = params.mainloop.ptr_cache_seqlens[idx_b]; + 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; @@ -356,7 +337,7 @@ class XeFMHAFwdKernel { 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_cache_seqlens[idx_b]; + 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); 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 015aab15d..595ad8fd0 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 @@ -64,10 +64,10 @@ void FmhaPrefillRunner<@HEAD_DIM@>::operator()(const Arguments& params) const { cute::Layout, cute::_1, cute::_1>, cute::Stride>; #if @APPENDKV_Q256_TILE@ - if (params.total_knew > 0 && params.seqlen_q == @APPENDKV_Q256_Q@ && - params.seqlen_knew == @APPENDKV_Q256_Q@ && + 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_knew == params.b * @APPENDKV_Q256_Q@ && !params.is_causal && !params.is_local && !use_sink && + 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>; 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 78a014d70..c94192cba 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 @@ -84,9 +84,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,13 +104,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__ cache_seqlens_old; + int* __restrict__ kv_cache_seqlens; // The stride between rows of Oaccum. int64_t oaccum_split_stride; @@ -124,16 +124,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; @@ -232,16 +232,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), }; @@ -252,7 +252,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; @@ -284,7 +291,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; } @@ -302,12 +309,12 @@ struct PrefillRunner { params.window_size_left, params.window_size_right}; if constexpr (CollectiveMainloop::AppendKV) { - mainloop_args.ptr_K_new = static_cast(params.knew_ptr); - mainloop_args.ptr_V_new = static_cast(params.vnew_ptr); - mainloop_args.ptr_cu_seqlens_k_new = params.cu_seqlens_knew; - mainloop_args.ptr_cache_seqlens = params.cache_seqlens_old; - mainloop_args.seq_len_kv_new = params.seqlen_knew; - mainloop_args.total_k_new = params.total_knew; + 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{ @@ -315,9 +322,9 @@ struct PrefillRunner { shape, static_cast(params.q_ptr), stride_Q, - CollectiveMainloop::DirectAppendKV ? static_cast(params.knew_ptr) : nullptr, + CollectiveMainloop::DirectAppendKV ? static_cast(params.k_new_ptr) : nullptr, stride_K, - CollectiveMainloop::DirectAppendKV ? static_cast(params.vnew_ptr) : nullptr, + CollectiveMainloop::DirectAppendKV ? static_cast(params.v_new_ptr) : nullptr, stride_V, static_cast(params.o_ptr), stride_O, @@ -486,10 +493,12 @@ struct FMHAConfig { 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_k > 0, "paged prefill requires positive max sequence lengths"); - TORCH_CHECK(params.total_q > 0 && params.total_k > 0, "paged prefill requires positive total sequence lengths"); - bool const has_append = params.total_knew > 0 && params.knew_ptr != nullptr && params.vnew_ptr != nullptr && - params.cache_seqlens_old != nullptr; + 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) { return run(params); } @@ -503,8 +512,8 @@ struct FMHAConfig { TORCH_CHECK(params.page_size > 0, "direct AppendKV requires a positive page_size"); TORCH_CHECK(params.max_num_pages_per_seq > 0, "direct AppendKV requires max_num_pages_per_seq"); TORCH_CHECK( - params.total_knew > 0 && params.knew_ptr != nullptr && params.vnew_ptr != nullptr && - params.cache_seqlens_old != nullptr, + params.total_kvnew > 0 && params.k_new_ptr != nullptr && params.v_new_ptr != nullptr && + params.kv_cache_seqlens != nullptr, "direct AppendKV requires k_new, v_new, and old cache lengths"); return run(params); } @@ -518,10 +527,12 @@ struct FMHAConfig { 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_k > 0, "non-paged prefill requires positive max sequence lengths"); - TORCH_CHECK(params.total_q > 0 && params.total_k > 0, "non-paged prefill requires positive total sequence lengths"); - bool const has_append = params.total_knew > 0 && params.knew_ptr != nullptr && params.vnew_ptr != nullptr && - params.cache_seqlens_old != nullptr; + 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(params); } From 0d94217fabe0fbaf23fb35f19e116d460bbaf85e Mon Sep 17 00:00:00 2001 From: Yuankun Shi Date: Wed, 29 Jul 2026 17:08:39 -0700 Subject: [PATCH 13/16] Consolidate prefill appendkv modes --- src/sycl/flash_attention.cpp | 20 +- .../collective/xe_fmha_fwd_mainloop.hpp | 308 ++++++++++-------- .../kernel/xe_fmha_fwd_kernel.hpp | 11 +- .../xe_fmha_fwd_prefill_kernel.cpp.in | 7 +- .../xe_fmha_fwd_prefill_runner.hpp | 57 ++-- 5 files changed, 226 insertions(+), 177 deletions(-) diff --git a/src/sycl/flash_attention.cpp b/src/sycl/flash_attention.cpp index 6a13fa6f6..e0f929c66 100644 --- a/src/sycl/flash_attention.cpp +++ b/src/sycl/flash_attention.cpp @@ -1007,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"); @@ -1113,10 +1122,10 @@ std::vector mha_fwd( params.window_size_right = window_size_right; params.total_q = total_q; params.total_kvcache = total_k; - params.b_k = batch_size_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; @@ -1462,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 @@ -1502,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) { 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 d952dd013..e074af2e7 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 @@ -55,6 +55,12 @@ using namespace cute; ///////////////////////////////////////////////////////////////////////////////////////////////// +enum class AppendKVMode { + kNone, + kStore, + kStoreAndDirectLoad, +}; + template < class DispatchPolicy_, bool CausalMask_, @@ -79,9 +85,7 @@ template < // 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 AppendKV_ = false, - bool DirectAppendKV_ = false, - bool WideAppendKV_ = false> + AppendKVMode AppendKVMode_ = AppendKVMode::kNone> struct FMHAFwdMainloop { static_assert(cutlass::detail::dependent_false, "Could not find a mainloop specialization."); }; @@ -108,9 +112,7 @@ template < class TiledCopyV_cache_, bool LocalMask_, bool PackGQA_, - bool AppendKV_, - bool DirectAppendKV_, - bool WideAppendKV_> + AppendKVMode AppendKVMode_> struct FMHAFwdMainloop< XeDefault, CausalMask_, @@ -131,9 +133,7 @@ struct FMHAFwdMainloop< TiledCopyV_cache_, LocalMask_, PackGQA_, - AppendKV_, - DirectAppendKV_, - WideAppendKV_> { + AppendKVMode_> { // // Type Aliases // @@ -211,12 +211,9 @@ 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 bool AppendKV = AppendKV_; - static constexpr bool DirectAppendKV = DirectAppendKV_; - static constexpr bool WideAppendKV = WideAppendKV_; - static_assert(!DirectAppendKV || (AppendKV && PagedKV), "Direct AppendKV requires paged AppendKV"); - static_assert( - !WideAppendKV || (AppendKV && PagedKV && !DirectAppendKV), "Wide AppendKV requires fused paged AppendKV"); + static constexpr AppendKVMode AppendMode = AppendKVMode_; + static constexpr bool AppendKV = AppendMode != AppendKVMode::kNone; + static constexpr bool AppendDirectLoad = AppendMode == AppendKVMode::kStoreAndDirectLoad && PagedKV; // User-facing arguments struct Arguments { @@ -314,6 +311,134 @@ struct FMHAFwdMainloop< 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, @@ -353,119 +478,38 @@ struct FMHAFwdMainloop< // 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) { - // The q256 HD64 specialization uses 32B vectors to reduce flattened - // store index overhead; other paths retain the 16B transaction. - constexpr int kVecElems = WideAppendKV ? 16 : 8; - using StoreVec = cute::conditional_t; - if (head_size_qk == head_size_vo && (head_size_qk % kVecElems) == 0) { - // 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) { + if constexpr (!AppendDirectLoad) { + 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; } - 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; - } - } - - 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; - } - } + } 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; } - return; } } @@ -624,15 +668,15 @@ struct FMHAFwdMainloop< 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); - if constexpr (DirectAppendKV) { + 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; - if ((cache_len_old % kTileKV) != 0 || (new_begin % kTileKV) != 0) { - barrier(); - } - } else { + direct_append = (cache_len_old % kTileKV) == 0 && (new_begin % kTileKV) == 0; + } + if (!direct_append) { barrier(); } } @@ -664,7 +708,7 @@ struct FMHAFwdMainloop< bool direct_append = false; int direct_block0 = kblocks_total; int direct_source_block0 = 0; - if constexpr (DirectAppendKV) { + 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; @@ -834,7 +878,7 @@ struct FMHAFwdMainloop< 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 (DirectAppendKV) { + 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_fwd_kernel.hpp b/src/sycl/kernels/flash_attention_v2/kernel/xe_fmha_fwd_kernel.hpp index 7df3fdcb8..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 @@ -102,7 +102,6 @@ class XeFMHAFwdKernel { using FragA = typename CollectiveMainloop::FragA; using FragARow = typename CollectiveMainloop::FragARow; - static constexpr bool DirectAppendKV = CollectiveMainloop::DirectAppendKV; // Tile scheduler derived types using TileScheduler = TileScheduler_; @@ -316,7 +315,7 @@ class XeFMHAFwdKernel { 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 (DirectAppendKV) { + 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] @@ -397,7 +396,9 @@ class XeFMHAFwdKernel { // 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; - if constexpr (DirectAppendKV) { + 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 @@ -411,8 +412,8 @@ class XeFMHAFwdKernel { 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(DirectAppendKV ? p.K : p.K_cache); - auto dcV = const_cast(DirectAppendKV ? p.V : p.V_cache); + 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); 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 595ad8fd0..7a9d44efe 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 @@ -77,13 +77,8 @@ void FmhaPrefillRunner<@HEAD_DIM@>::operator()(const Arguments& params) const { using AppendSubgroupLayoutQK = cute::Layout< cute::Shape, cute::_1, cute::_1>, cute::Stride>; -#if @HEAD_DIM@ == 64 - FMHAConfig::run_paged_wide_append(params); -#else FMHAConfig::run_paged_direct_append(params); -#endif + AppendSubgroupLayoutQK>::run_paged(params); return; } #endif 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 c94192cba..65e5525e5 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 @@ -322,9 +322,9 @@ struct PrefillRunner { shape, static_cast(params.q_ptr), stride_Q, - CollectiveMainloop::DirectAppendKV ? static_cast(params.k_new_ptr) : nullptr, + CollectiveMainloop::AppendKV ? static_cast(params.k_new_ptr) : nullptr, stride_K, - CollectiveMainloop::DirectAppendKV ? static_cast(params.v_new_ptr) : nullptr, + CollectiveMainloop::AppendKV ? static_cast(params.v_new_ptr) : nullptr, stride_V, static_cast(params.o_ptr), stride_O, @@ -402,10 +402,8 @@ struct FMHAConfig { bool isVarLen, bool CachedKV, bool PagedKV, - bool AppendKV, - class Scheduler, - bool DirectAppendKV = false, - bool WideAppendKV = false> + 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. @@ -457,9 +455,7 @@ struct FMHAConfig { GmemTiledCopyV_cache, LocalMask, false, - AppendKV, - DirectAppendKV, - WideAppendKV>; + AppendMode>; // Epilogue using CollectiveEpilogue = @@ -500,26 +496,17 @@ struct FMHAConfig { 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(params); + constexpr auto append_mode = int(get<1>(TileShapeOutput{})) == 64 + ? cutlass::fmha::collective::AppendKVMode::kStore + : cutlass::fmha::collective::AppendKVMode::kStoreAndDirectLoad; + return run(params); } - return run(params); - } - - static int run_paged_direct_append(const Arguments& params) { - TORCH_CHECK(params.cu_seqlens_q != nullptr, "direct AppendKV requires cu_seqlens_q"); - TORCH_CHECK(params.cu_seqlens_k != nullptr, "direct AppendKV requires per-batch cache lengths"); - TORCH_CHECK(params.page_table != nullptr, "direct AppendKV requires page_table"); - TORCH_CHECK(params.page_size > 0, "direct AppendKV requires a positive page_size"); - TORCH_CHECK(params.max_num_pages_per_seq > 0, "direct AppendKV requires max_num_pages_per_seq"); - TORCH_CHECK( - params.total_kvnew > 0 && params.k_new_ptr != nullptr && params.v_new_ptr != nullptr && - params.kv_cache_seqlens != nullptr, - "direct AppendKV requires k_new, v_new, and old cache lengths"); - return run(params); - } - - static int run_paged_wide_append(const Arguments& 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. @@ -534,9 +521,19 @@ struct FMHAConfig { 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(params); + return run< + true, + true, + false, + cutlass::fmha::collective::AppendKVMode::kStore, + cutlass::fmha::kernel::XeFHMAIndividualTileScheduler>(params); } - return run(params); + return run< + true, + true, + false, + cutlass::fmha::collective::AppendKVMode::kNone, + cutlass::fmha::kernel::XeFHMAIndividualTileScheduler>(params); } static int run(const Arguments& params) { From 1c4b521a9cad9faac4a695af12c9e652adeccb9f Mon Sep 17 00:00:00 2001 From: Yuankun Shi Date: Thu, 30 Jul 2026 10:47:24 -0700 Subject: [PATCH 14/16] Add wide AppendKV enum mode --- .../flash_attention_v2/collective/xe_fmha_fwd_mainloop.hpp | 4 +++- .../kernels/flash_attention_v2/xe_fmha_fwd_prefill_runner.hpp | 2 +- 2 files changed, 4 insertions(+), 2 deletions(-) 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 e074af2e7..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 @@ -58,6 +58,7 @@ using namespace cute; enum class AppendKVMode { kNone, kStore, + kStoreWide, kStoreAndDirectLoad, }; @@ -214,6 +215,7 @@ struct FMHAFwdMainloop< 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 { @@ -478,7 +480,7 @@ struct FMHAFwdMainloop< // 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 (!AppendDirectLoad) { + if constexpr (AppendWideStore) { if (store_kv_new_vectorized<16, cutlass::ulonglong4>( K_cache_2D, V_cache_2D, 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 65e5525e5..e263d9749 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 @@ -497,7 +497,7 @@ struct FMHAConfig { params.kv_cache_seqlens != nullptr; if (has_append) { constexpr auto append_mode = int(get<1>(TileShapeOutput{})) == 64 - ? cutlass::fmha::collective::AppendKVMode::kStore + ? cutlass::fmha::collective::AppendKVMode::kStoreWide : cutlass::fmha::collective::AppendKVMode::kStoreAndDirectLoad; return run(params); } From b1fe867c332ad314b543f80d5f14bcd1428f0254 Mon Sep 17 00:00:00 2001 From: Yuankun Shi Date: Thu, 30 Jul 2026 12:38:41 -0700 Subject: [PATCH 15/16] Restore AppendKV mode dispatch --- .../flash_attention_v2/xe_fmha_fwd_prefill_kernel.cpp.in | 4 +++- .../flash_attention_v2/xe_fmha_fwd_prefill_runner.hpp | 9 +++++---- 2 files changed, 8 insertions(+), 5 deletions(-) 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 7a9d44efe..704ef1711 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 @@ -77,8 +77,10 @@ void FmhaPrefillRunner<@HEAD_DIM@>::operator()(const Arguments& params) const { 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::run_paged(params); + AppendSubgroupLayoutQK>::template run_paged(params); return; } #endif 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 e263d9749..99f07bcb1 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 @@ -483,6 +483,10 @@ 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) { 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"); @@ -496,10 +500,7 @@ struct FMHAConfig { 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) { - constexpr auto append_mode = int(get<1>(TileShapeOutput{})) == 64 - ? cutlass::fmha::collective::AppendKVMode::kStoreWide - : cutlass::fmha::collective::AppendKVMode::kStoreAndDirectLoad; - return run(params); + return run(params); } return run< true, From 67e48f6deea0cebaf904daf0153b1627d953db91 Mon Sep 17 00:00:00 2001 From: Yuankun Shi Date: Thu, 6 Aug 2026 23:01:47 -0700 Subject: [PATCH 16/16] Use a standalone prepass for large AppendKV prefill --- src/FMHAPrefillXe20.cmake | 15 ++ .../kernel/xe_fmha_append_kv.hpp | 156 ++++++++++++++++++ .../xe_fmha_fwd_prefill_kernel.cpp.in | 23 +++ .../xe_fmha_fwd_prefill_runner.hpp | 45 +++++ 4 files changed, 239 insertions(+) create mode 100644 src/sycl/kernels/flash_attention_v2/kernel/xe_fmha_append_kv.hpp diff --git a/src/FMHAPrefillXe20.cmake b/src/FMHAPrefillXe20.cmake index b5f76f209..ed0315648 100644 --- a/src/FMHAPrefillXe20.cmake +++ b/src/FMHAPrefillXe20.cmake @@ -49,6 +49,11 @@ set(FMHA_PREFILL_NUM_SG_192 16) set(FMHA_PREFILL_TILED_Q_256 128) set(FMHA_PREFILL_TILED_KV_256 64) 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) @@ -132,6 +137,16 @@ foreach(HEAD_DIM ${FMHA_PREFILL_PAGED_HEAD_DIMS}) 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) 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/xe_fmha_fwd_prefill_kernel.cpp.in b/src/sycl/kernels/flash_attention_v2/xe_fmha_fwd_prefill_kernel.cpp.in index 704ef1711..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 @@ -98,6 +98,29 @@ void FmhaPrefillRunner<@HEAD_DIM@>::operator()(const Arguments& params) const { } #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 && 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 99f07bcb1..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" @@ -193,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; @@ -500,6 +533,18 @@ struct FMHAConfig { 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<