diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 2aa9cb3cafb7d..acaa3b724fc51 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -4532,7 +4532,7 @@ This version of the operator has been available since version 1 of the 'com.micr
v_scale (optional) : T_KV_SCALE
Dequantization scale of the value cache. Shape is (1) when 'v_quant_type' is 'PER_TENSOR' and (kv_num_heads, 1, head_size) when it is 'PER_CHANNEL'. Quantization is symmetric (no zero point).
attention_metadata (optional) : S
-
1D tensor with shape (2) holding [max_query_len_bound, max_kv_len_bound] in CPU memory. max_query_len_bound is an upper bound on the number of new tokens any one sequence contributes; max_kv_len_bound is an upper bound on past_seqlens[i] + query_len[i]. Both are replay-wide upper bounds, never exact per-step values: they must hold for every step this node -- or a CUDA Graph capturing it -- will serve, and 0 means 'unknown'. They may only select the backend and size launch dimensions and workspaces; they never enter a mask comparison, so over-estimating only costs empty work. The op can otherwise obtain these only by copying 'cumulative_sequence_length' and 'past_seqlens' back from the device and synchronizing the stream on every call, which stalls the pipeline once per node per step and makes the op impossible to capture into a CUDA Graph. Schedulers already track these bounds on the host, so supplying them is normally free. When absent, the op falls back to the device readback. The values are trusted: an under-sized bound violates the contract and may omit attention work.
+
1D tensor with shape (2) or (3) holding [max_query_len_bound, max_kv_len_bound, optional max_kv_len_lower_bound] in CPU memory. max_query_len_bound is an upper bound on the number of new tokens any one sequence contributes; max_kv_len_bound is an upper bound on past_seqlens[i] + query_len[i]. Both are replay-wide upper bounds, never exact per-step values: they must hold for every step this node -- or a CUDA Graph capturing it -- will serve, and 0 means 'unknown'. They may only select the backend and size launch dimensions and workspaces; they never enter a mask comparison, so over-estimating only costs empty work. max_kv_len_lower_bound is a replay-wide lower bound on the largest per-sequence KV length in the batch and 0 means 'unknown'. It is a provider-neutral performance hint; omitting it preserves the shape-(2) contract and disables optimizations that require a lower bound unless the op reads exact lengths back from the device. The op can otherwise obtain the upper bounds only by copying 'cumulative_sequence_length' and 'past_seqlens' back from the device and synchronizing the stream on every call, which stalls the pipeline once per node per step and makes the op impossible to capture into a CUDA Graph. Schedulers already track these bounds on the host, so supplying them is normally free. When absent, the op falls back to the device readback. The upper bounds are trusted: an under-sized bound violates the contract and may omit attention work.
#### Outputs (1 - 3) @@ -7286,5 +7286,3 @@ No versioning maintained for experimental ops.
T : tensor(float)
Constrain input and output types to float32 tensors.
- - diff --git a/docs/contrib_ops/cuda/paged_attention.md b/docs/contrib_ops/cuda/paged_attention.md index 3619a936ed71e..1d5f2d78b316c 100644 --- a/docs/contrib_ops/cuda/paged_attention.md +++ b/docs/contrib_ops/cuda/paged_attention.md @@ -174,7 +174,7 @@ matches the landing order in [§19](#19-phasing), so the schema grows monotonica | 13 | `k_norm_weight` | `T` (opt) | `(head_size,)` | **new — §7** | | 14 | `k_scale` | `T_KV_SCALE` (opt) | `(1,)` or `(kv_num_heads, 1, head_size)` | **new — §8** | | 15 | `v_scale` | `T_KV_SCALE` (opt) | `(1,)` or `(kv_num_heads, 1, head_size)` | **new — §8**; absent in `LATENT` | -| 16 | `attention_metadata` | `S` (opt, **CPU**) | `(2,)` | **new — trusted bounds only, §4.7** | +| 16 | `attention_metadata` | `S` (opt, **CPU**) | `(2,)` or `(3,)` | **new — trusted bounds only, §4.7** | | 17 | `query_positions` | `S` (opt) | `(token_count,)` | **new — §4.8** | | 18 | `attention_bias` | `T` (opt) | `(batch_size or 1, num_heads or 1, query_length_capacity, context_length_capacity)` | **new — §10** | @@ -316,7 +316,8 @@ Three derivations satisfy the rule, none of which needs a synchronization: | Quantity | Source | Replay-safe because | |---|---|---| | decode vs. prefill dispatch | static shapes: `query.shape[0] == cumulative_sequence_length.shape[0] - 1` | shapes are fixed for a captured graph | -| grid size, split count, gather/workspace extents | static capacity bound `max_kv_len_bound = block_table.shape[1] * block_size` | independent of step | +| grid size, gather/workspace extents | static capacity bound `max_kv_len_bound = block_table.shape[1] * block_size` | independent of step | +| split-KV eligibility | `max_kv_len_lower_bound` lower bound, when supplied | proves splitting is worthwhile on every replay | | per-sequence KV length, causal and window masking, gather trip counts | device `past_seqlens` / `cumulative_sequence_length` | re-read from device memory on every replay | The shape test is a **performance heuristic only**. `token_count <= batch_size` does not prove that @@ -333,12 +334,14 @@ per-step sync. `attention_metadata` is consequently demoted to optional **replay-wide bounds**: ```text -attention_metadata : (2,) int32, OrtMemTypeCPUInput +attention_metadata : (2,) or (3,) int32, OrtMemTypeCPUInput [0] max_query_len_bound # 0 = unknown. Replay-wide upper bound on tokens from any one sequence. [1] max_kv_len_bound # 0 = unknown. Replay-wide upper bound on total KV length of any sequence. + [2] max_kv_len_lower_bound # Optional; 0 = unknown. Replay-wide lower bound on the largest + # per-sequence KV length in the batch. ``` -- Both entries are **upper bounds, never exact values**, and must hold for *every* step the node — +- The first two entries are **upper bounds, never exact values**, and must hold for *every* step the node — or the captured graph containing it — will serve. - `0` means "no bound"; the implementation falls back to `token_count` for query length and to `block_table.shape[1] * block_size` for KV length. @@ -349,6 +352,13 @@ attention_metadata : (2,) int32, OrtMemTypeCPUInput - A valid bound may only shrink launch dimensions and workspace sizes. It must not enter a mask comparison. Device loops use the current device lengths, additionally bounded by the trusted launch/workspace extent. +- `max_kv_len_lower_bound` is a provider-neutral, performance-only lower bound on + `max_i(past_seqlens[i] + query_len[i])`. CUDA currently uses it for split-KV eligibility, but its + contract does not name or require that backend. Shape `(2,)` remains supported and defaults this + value to `0` (unknown), disabling split-KV unless an exact device readback is already required. +- The upper bound or block-table capacity cannot be substituted for this lower bound. A producer may + pad a short live sequence to a large replay capacity, and treating that capacity as live length + would force split/combine overhead on every early replay. - Even for an invalid bound, every device read must remain memory-safe: device lengths and static tensor capacities guard all accesses. This safety property does not imply a correct result when the producer violates the upper-bound contract. @@ -376,7 +386,8 @@ which must now be sized by `batch_size * max_kv_len_bound` so that the allocatio > | Host quantity | Consumer | Bound used | > |---|---|---| > | `max_query_len` | `params.seqlen_q` (Flash), `p.sequence_length` (MEA) — grid extent only | `max_query_len_bound`, else `token_count` | -> | `max_kv_len` | quantized-Flash `max_seqlen_k`, decode split count | `max_kv_len_bound`, else `block_table.shape[1] * block_size` | +> | `max_kv_len` | quantized-Flash `max_seqlen_k`, split workspace sizing | `max_kv_len_bound`, else `block_table.shape[1] * block_size` | +> | split-KV eligibility | FlashAttention decode dispatch | `max_kv_len_lower_bound`, else disabled unless exact lengths were read back | > | `total_kv_tokens` | gather staging buffer extent | `batch_size * max_kv_len_bound` | > > Two narrow cases still take the readback, and only when the caller supplied **no** metadata at all: @@ -799,8 +810,10 @@ Mirror GQA: `onnxruntime_USE_FP8_KV_CACHE` (default ON), `onnxruntime_USE_INT4_K > full prefill) and a wrong heuristic only costs speed. That is what removes the D→H sync. > - Split-KV: `ComputePagedDecodeSplits` splits the KV range across up to 32 CTAs only when > `token_count * num_heads` would leave the device under-occupied. `max_kv_len` may be an upper -> bound. Empty splits publish `(max = -FLT_MAX, denom = 0)` and the reduce kernel skips them, so -> their accumulator slice is never read. +> bound. FlashAttention keeps the replay-wide split count and workspaces fixed while partitioning +> each varlen sequence from its live device length, so loose upper bounds do not concentrate useful +> tiles in the first split. Empty splits publish `(max = -FLT_MAX, denom = 0)` and the reduce kernel +> skips them, so their accumulator slice is never read. > - The `FlashAttention` / `EfficientAttention` prologue (packed-QKV unpack, fused QK-Norm + rotary, > `ReshapeAndCache`) was factored into a shared `PrepareQueryAndCache`, which the decode backend > reuses. It lives outside the `USE_FLASH_ATTENTION` / `USE_MEMORY_EFFICIENT_ATTENTION` guards @@ -1351,10 +1364,10 @@ Consolidated, to be implemented in `paged_attention_helper::CheckInputs`. Every every logical element type this operator stores is expressible as an ONNX element type. The reserved sub-byte values are rejected until a `uint8` packed cache exists. FP8 availability is controlled by `onnxruntime_USE_FP8_KV_CACHE`, without an additional runtime architecture gate. -- `attention_metadata`: rank 1, `dim0 == 2`, `int32`, CPU-resident; entries `>= 0`; each non-zero - entry is clamped to its static limit before use and may only size launch dimensions and workspace - (§4.7). It must never enter a mask comparison. Each value is a trusted upper bound for every step - served by the node or captured graph. +- `attention_metadata`: rank 1, `dim0 ∈ {2, 3}`, `int32`, CPU-resident; entries `>= 0`; the first + two entries are trusted upper bounds and the optional third is a trusted lower bound for every + step served by the node or captured graph (§4.7). Bounds may only select implementations or size + launch dimensions and workspace; they must never enter a mask comparison. - `query_positions`: rank 1, `dim0 == token_count`, `int32`, entries `>= 0`. - `attention_bias`: rank 4; `dim0 ∈ {1, batch_size}`, `dim1 ∈ {1, num_heads}`, `dim2 == query_length_capacity`, and `dim3 == context_length_capacity`; both capacities must cover diff --git a/onnxruntime/contrib_ops/cpu/bert/paged_attention_helper.h b/onnxruntime/contrib_ops/cpu/bert/paged_attention_helper.h index a1ed3dfa362e7..2795dfb1e6220 100644 --- a/onnxruntime/contrib_ops/cpu/bert/paged_attention_helper.h +++ b/onnxruntime/contrib_ops/cpu/bert/paged_attention_helper.h @@ -568,15 +568,15 @@ Status CheckInputs(const T* query, ORT_RETURN_IF_ERROR(CheckKVCacheDataType(k_cache_dtype, cache_storage_dtype, "k_cache_dtype")); ORT_RETURN_IF_ERROR(CheckKVCacheDataType(v_cache_dtype, cache_storage_dtype, "v_cache_dtype")); - // Optional host-side [max_query_len_bound, max_kv_len_bound]. Only the shape is checked here. - // The entries are *trusted upper bounds* and cannot be cross-checked against the device tensors - // they bound without the readback this input exists to remove; see the trust boundary in - // docs/contrib_ops/cuda/paged_attention.md section 4.7. + // Optional host-side [max_query_len_bound, max_kv_len_bound, max_kv_len_lower_bound]. + // The first two entries are trusted upper bounds and cannot be cross-checked against the device + // tensors they bound without the readback this input exists to remove. The optional third entry + // is a performance-only lower bound. See docs/contrib_ops/cuda/paged_attention.md section 4.7. if (attention_metadata != nullptr) { const auto& metadata_dims = attention_metadata->Shape().GetDims(); - if (metadata_dims.size() != 1 || metadata_dims[0] != 2) { + if (metadata_dims.size() != 1 || (metadata_dims[0] != 2 && metadata_dims[0] != 3)) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Input 'attention_metadata' must have shape (2), got ", + "Input 'attention_metadata' must have shape (2) or (3), got ", attention_metadata->Shape().ToString()); } } diff --git a/onnxruntime/contrib_ops/cuda/bert/attention_data.h b/onnxruntime/contrib_ops/cuda/bert/attention_data.h index 98782b7c88290..6d8448b0eb65f 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention_data.h +++ b/onnxruntime/contrib_ops/cuda/bert/attention_data.h @@ -284,10 +284,11 @@ struct PagedAttentionData { const T* q_norm_weight = nullptr; const T* k_norm_weight = nullptr; - // Flash buffers. FlashAttention always emits FP32 log-sum-exp regardless of T; with - // params.num_splits <= 1 (which mha_varlen_fwd never overrides) the varlen layout is - // [num_heads, token_count]. + // Flash buffers. FlashAttention always emits FP32 log-sum-exp regardless of T. float* softmax_lse = nullptr; + float* flash_softmax_lse_accum = nullptr; + float* flash_out_accum = nullptr; + int flash_num_splits = 0; int* cumulative_seqlens_kv = nullptr; // Flash api takes cumulative sequence length for kv-cache // Fused op buffers diff --git a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_api.cc b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_api.cc index d2424ea526a5b..6a96039c70943 100644 --- a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_api.cc +++ b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_api.cc @@ -363,7 +363,10 @@ Status mha_varlen_fwd(const cudaDeviceProp& dprops, bool is_bf16, int local_window_size, int max_num_blocks_per_seq, - int page_block_size) { + int page_block_size, + int num_splits, + void* softmax_lse_accum, + void* out_accum) { auto round_multiple = [](int x, int m) { return (x + m - 1) / m * m; }; const int head_size_rounded = round_multiple(head_size, 32); const int seqlen_q_rounded = round_multiple(max_seqlen_q, 128); @@ -397,6 +400,20 @@ Status mha_varlen_fwd(const cudaDeviceProp& dprops, params.total_q = total_q; params.dprops = &dprops; + const bool pure_decode = max_seqlen_q == 1 && total_q == batch_size; + if (num_splits > 1) { + ORT_RETURN_IF_NOT(pure_decode, + "FlashAttention varlen split-KV requires exactly one query token per sequence."); + ORT_RETURN_IF_NOT(softmax_lse_accum != nullptr && out_accum != nullptr, + "FlashAttention varlen split-KV requires LSE and output accumulator workspaces."); + // Varlen normally uses cu_seqlens_q and leaves the batch stride at zero. In pure decode the + // packed output has exactly one row per batch, so the split-combine kernel can address it as + // a dense [batch, 1, heads, head_size] tensor. + params.o_batch_stride = num_heads * head_size; + params.num_splits = num_splits; + params.softmax_lseaccum_ptr = softmax_lse_accum; + params.oaccum_ptr = out_accum; + } if (paged_KV) { params.block_table = block_table; params.block_table_batch_stride = max_num_blocks_per_seq; diff --git a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_api.h b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_api.h index 32ea803c998d8..a72c1ea2331aa 100644 --- a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_api.h +++ b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_api.h @@ -90,7 +90,10 @@ Status mha_varlen_fwd(const cudaDeviceProp& dprops, bool is_bf16, int local_window_size = -1, int max_num_blocks_per_seq = 0, - int page_block_size = 1); + int page_block_size = 1, + int num_splits = 0, + void* softmax_lse_accum = nullptr, + void* out_accum = nullptr); Status mha_fwd_kvcache(const cudaDeviceProp& dprops, cudaStream_t stream, diff --git a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_kernel.h b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_kernel.h index e80d632c3b77a..a17fb92bd300e 100644 --- a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_kernel.h +++ b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_fwd_kernel.h @@ -478,7 +478,12 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params& params, cons // if (threadIdx.x == 0 && blockIdx.y == 1 && blockIdx.z == 0) { printf("params.knew_ptr = %p, seqlen_k_cache + seqlen_knew = %d\n", params.knew_ptr, binfo.seqlen_k_cache + (params.knew_ptr == nullptr ? 0 : params.seqlen_knew)); } if (m_block * kBlockM >= binfo.actual_seqlen_q) return; - const int n_blocks_per_split = ((params.seqlen_k + kBlockN - 1) / kBlockN + num_n_splits - 1) / num_n_splits; + // CUDA graph replay keeps the split count and workspace sizes fixed, but varlen sequence lengths + // remain live device values. Partition the live tiles so a loose replay upper bound does not + // concentrate all useful work in the first split. Fixed-length callers retain the original extent. + const int n_blocks_to_split = + cute::ceil_div(Is_even_MN ? params.seqlen_k : binfo.actual_seqlen_k, kBlockN); + const int n_blocks_per_split = cute::ceil_div(n_blocks_to_split, num_n_splits); const int n_block_min = !Is_local ? n_split_idx * n_blocks_per_split : std::max(n_split_idx * n_blocks_per_split, (m_block * kBlockM + binfo.actual_seqlen_k - binfo.actual_seqlen_q - params.window_size_left) / kBlockN); diff --git a/onnxruntime/contrib_ops/cuda/bert/paged_attention.cc b/onnxruntime/contrib_ops/cuda/bert/paged_attention.cc index 3e3f4c7c96c10..51e7bbf796916 100644 --- a/onnxruntime/contrib_ops/cuda/bert/paged_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/paged_attention.cc @@ -4,6 +4,7 @@ #include #include #include +#include #include "core/providers/cuda/cuda_common.h" #include "core/platform/env_var_utils.h" @@ -25,6 +26,8 @@ namespace onnxruntime { namespace contrib { namespace cuda { +constexpr int kFlashSplitKvMinSequenceLength = 512; + #define REGISTER_KERNEL_TYPED(T, TCACHE) \ ONNX_OPERATOR_TYPED_KERNEL_EX( \ PagedAttention, \ @@ -375,9 +378,10 @@ Status PagedAttention::ComputeInternal(OpKernelContext* context) cons IAllocatorUniquePtr gathered_value_buffer; IAllocatorUniquePtr fmha_buffer; - // 'attention_metadata' supplies replay-wide *upper bounds* on the per-sequence query and KV - // lengths (docs/contrib_ops/cuda/paged_attention.md section 4.7). Bounds are all the backends - // need from the host: they only select the kernel, size launch dimensions and size workspaces. + // 'attention_metadata' supplies replay-wide upper bounds on the per-sequence query and KV + // lengths, plus an optional replay-wide lower bound on the largest KV length + // (docs/contrib_ops/cuda/paged_attention.md section 4.7). Bounds are all the backends need from + // the host: they only select the kernel, size launch dimensions and size workspaces. // Every per-sequence length that enters a mask is re-read from device memory by the kernel // itself, which is what keeps a captured graph correct as the sequences grow. // @@ -389,14 +393,18 @@ Status PagedAttention::ComputeInternal(OpKernelContext* context) cons const bool has_metadata_bounds = attention_metadata != nullptr; int max_query_len_bound = parameters.token_count; int max_kv_len_bound = max_kv_len_capacity; + int max_kv_len_lower_bound = 0; if (has_metadata_bounds) { const int* metadata = attention_metadata->Data(); const int metadata_query_bound = metadata[0]; const int metadata_kv_bound = metadata[1]; - if (metadata_query_bound < 0 || metadata_kv_bound < 0) { + const int metadata_kv_lower_bound = + attention_metadata->Shape().Size() == 3 ? metadata[2] : 0; + if (metadata_query_bound < 0 || metadata_kv_bound < 0 || metadata_kv_lower_bound < 0) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "PagedAttention: 'attention_metadata' entries must be non-negative, got [", - metadata_query_bound, ", ", metadata_kv_bound, "]. Use 0 for 'unknown'."); + metadata_query_bound, ", ", metadata_kv_bound, ", ", + metadata_kv_lower_bound, "]. Use 0 for 'unknown'."); } // Clamp each bound to the static limit it can never exceed, so an over-large (or unknown) // bound degrades to the same sizing we would use with no metadata at all. @@ -406,6 +414,13 @@ Status PagedAttention::ComputeInternal(OpKernelContext* context) cons if (metadata_kv_bound > 0 && metadata_kv_bound < max_kv_len_bound) { max_kv_len_bound = metadata_kv_bound; } + if (metadata_kv_lower_bound > max_kv_len_bound) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "PagedAttention: attention_metadata max_kv_len_lower_bound (", + metadata_kv_lower_bound, ") must not exceed max_kv_len_bound (", + max_kv_len_bound, ")."); + } + max_kv_len_lower_bound = metadata_kv_lower_bound; } // Backend selection from static shapes alone. @@ -510,6 +525,7 @@ Status PagedAttention::ComputeInternal(OpKernelContext* context) cons } } total_kv_tokens = cum_kv_pinned.get()[parameters.batch_size]; + max_kv_len_lower_bound = max_kv_len; if (total_kv_tokens <= 0) { return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "PagedAttention: total_kv_tokens is not positive (", total_kv_tokens, @@ -583,6 +599,34 @@ Status PagedAttention::ComputeInternal(OpKernelContext* context) cons #endif auto softmax_lse_buffer = GetScratchBuffer(softmax_lse_bytes, GetComputeStream(context)); + int flash_num_splits = 0; + IAllocatorUniquePtr flash_softmax_lse_accum_buffer; + IAllocatorUniquePtr flash_out_accum_buffer; +#if USE_FLASH_ATTENTION + const bool use_flash_split_kv = + use_flash_attention && + parameters.token_count == parameters.batch_size && + max_query_len == 1 && + !parameters.use_smooth_softmax && + parameters.local_window_size <= 0; + if (use_flash_split_kv) { + // The combine kernel costs more than it saves for short decode contexts, even when a high + // query-head count makes the occupancy heuristic select two splits. The upper bound sizes the + // workspaces, while the replay-wide lower bound proves splitting is worthwhile on every replay. + if (max_kv_len_lower_bound > kFlashSplitKvMinSequenceLength) { + const auto [num_splits, softmax_lse_accum_bytes, out_accum_bytes] = + onnxruntime::flash::get_num_splits_and_buffer_sizes( + parameters.batch_size, 1, max_kv_len, parameters.num_heads, + parameters.head_size, device_prop.multiProcessorCount); + flash_num_splits = static_cast(num_splits); + flash_softmax_lse_accum_buffer = + GetScratchBuffer(softmax_lse_accum_bytes, GetComputeStream(context)); + flash_out_accum_buffer = + GetScratchBuffer(out_accum_bytes, GetComputeStream(context)); + } + } +#endif + if (needs_dense_kv) { const size_t gather_elems = static_cast(total_kv_tokens) * gathered_num_heads * parameters.head_size; @@ -662,7 +706,9 @@ Status PagedAttention::ComputeInternal(OpKernelContext* context) cons debug_info.use_flash_attention = use_flash_attention; debug_info.use_efficient_attention = use_memory_efficient_attention; debug_info.use_decoder_attention = use_paged_decode && !use_xqa_decode; - if (use_paged_decode && !use_xqa_decode) { + if (use_flash_attention) { + debug_info.num_splits = std::max(1, flash_num_splits); + } else if (use_paged_decode && !use_xqa_decode) { debug_info.num_splits = num_splits; } debug_info.gqa_group_size = parameters.num_heads / parameters.kv_num_heads; @@ -706,6 +752,10 @@ Status PagedAttention::ComputeInternal(OpKernelContext* context) cons // FlashAttention always writes fp32 log-sum-exp, independent of T. data.softmax_lse = reinterpret_cast(softmax_lse_buffer.get()); } + data.flash_num_splits = flash_num_splits; + data.flash_softmax_lse_accum = reinterpret_cast(flash_softmax_lse_accum_buffer.get()); + data.flash_out_accum = reinterpret_cast(flash_out_accum_buffer.get()); + data.max_kv_len = max_kv_len; if (workspace_buffer != nullptr) { data.workspace_buffer = reinterpret_cast(workspace_buffer.get()); } @@ -718,7 +768,6 @@ Status PagedAttention::ComputeInternal(OpKernelContext* context) cons data.gathered_key = reinterpret_cast(gathered_key_buffer.get()); data.gathered_value = reinterpret_cast(gathered_value_buffer.get()); data.total_kv_tokens = total_kv_tokens; - data.max_kv_len = max_kv_len; } if (use_paged_decode && !use_xqa_decode) { data.decode_partial_out = reinterpret_cast(decode_partial_out_buffer.get()); diff --git a/onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.cu index 2314fe0c9e74d..e593ea6df9c30 100644 --- a/onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.cu @@ -1656,9 +1656,7 @@ Status FlashAttention( if constexpr (IsQuantizedCache::value) { // FlashAttention cannot read a quantized page, so dequantize the live context into a dense // packed-varlen [total_kv_tokens, kv_num_heads, head_size] buffer (no GQA expansion — Flash - // does the grouping itself) and use the non-paged varlen entry point. That path leaves - // params.num_splits at 0 exactly like the paged one, so the fp32 [num_heads, token_count] - // softmax_lse layout the head-sink epilogue relies on is unchanged. + // does the grouping itself) and use the non-paged varlen entry point. ORT_RETURN_IF_ERROR((LaunchGatherAndExpandPagedKVCache( data.key_cache, data.value_cache, data.gathered_key, data.gathered_value, data.k_scale, data.v_scale, k_per_channel, v_per_channel, @@ -1670,23 +1668,22 @@ Status FlashAttention( reinterpret_cast(data.gathered_value), output, cumulative_seqlens_q, cumulative_seqlens_kv, /*seqused_k*/ nullptr, /*block_table*/ nullptr, softmax_lse, batch_size, num_heads, kv_num_heads, head_size, max_query_len, data.max_kv_len, token_count, scale, softcap, /*is_causal*/ true, is_bf16, - local_window_size - 1)); + local_window_size - 1, /*max_num_blocks_per_seq*/ 0, /*page_block_size*/ 1, + data.flash_num_splits, data.flash_softmax_lse_accum, data.flash_out_accum)); } else { void* key_cache = reinterpret_cast(data.key_cache); void* value_cache = reinterpret_cast(data.value_cache); - const int max_seq_len = max_num_blocks_per_seq * block_size; ORT_RETURN_IF_ERROR(onnxruntime::flash::mha_varlen_fwd( device_prop, stream, q, key_cache, value_cache, output, cumulative_seqlens_q, cumulative_seqlens_kv, /*seqused_k*/ nullptr, block_table, softmax_lse, batch_size, num_heads, kv_num_heads, head_size, - max_query_len, max_seq_len, token_count, scale, softcap, /*is_causal*/ true, is_bf16, local_window_size - 1, - max_num_blocks_per_seq, block_size)); + max_query_len, data.max_kv_len, token_count, scale, softcap, /*is_causal*/ true, is_bf16, + local_window_size - 1, + max_num_blocks_per_seq, block_size, + data.flash_num_splits, data.flash_softmax_lse_accum, data.flash_out_accum)); } if (parameters.use_smooth_softmax) { - // Rescale by the softmax denominator that the sink logit adds. mha_varlen_fwd leaves - // params.num_splits at 0, so the split-combine kernel never runs and softmax_lse carries the - // unpadded [num_heads, token_count] fp32 layout this epilogue expects. If varlen ever enables - // num_splits > 1, both the layout and this epilogue must be revisited. + // Sink-bearing steps remain unsplit until the split-combine LSE layout is qualified. ORT_RETURN_IF_ERROR(LaunchApplyHeadSink(data.output, data.softmax_lse, data.head_sink, token_count, num_heads, head_size, stream, max_threads_per_block)); } diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 2797ce8f4872f..d53ee2ace22cf 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -1434,7 +1434,9 @@ cumulative_sequence_length records cumulated length of each sequence length. // Input 'k_norm_weight': (head_size) // Input 'k_scale': (1) for PER_TENSOR, (kv_num_heads, 1, head_size) for PER_CHANNEL // Input 'v_scale': (1) for PER_TENSOR, (kv_num_heads, 1, head_size) for PER_CHANNEL -// Input 'attention_metadata': (2), CPU memory: [max_query_len_bound, max_kv_len_bound] +// Input 'attention_metadata': (2) or (3), CPU memory: +// [max_query_len_bound, max_kv_len_bound, +// optional max_kv_len_lower_bound] // Output 'output': (token_count, num_heads * v_head_size) // Output 'key_cache_out': (num_blocks, block_size, kv_num_heads, head_size) // Output 'value_cache_out': (num_blocks, block_size, kv_num_heads, head_size), absent for LATENT @@ -1718,18 +1720,24 @@ ONNX_MS_OPERATOR_SET_SCHEMA( OpSchema::Optional) .Input(16, "attention_metadata", - "1D tensor with shape (2) holding [max_query_len_bound, max_kv_len_bound] in CPU memory. " + "1D tensor with shape (2) or (3) holding [max_query_len_bound, max_kv_len_bound, " + "optional max_kv_len_lower_bound] in CPU memory. " "max_query_len_bound is an upper bound on the number of new tokens any one sequence " "contributes; max_kv_len_bound is an upper bound on past_seqlens[i] + query_len[i]. Both are " "replay-wide upper bounds, never exact per-step values: they must hold for every step this node " "-- or a CUDA Graph capturing it -- will serve, and 0 means 'unknown'. They may only select the " "backend and size launch dimensions and workspaces; they never enter a mask comparison, so " - "over-estimating only costs empty work. The op can otherwise obtain these only by copying " + "over-estimating only costs empty work. max_kv_len_lower_bound is a replay-wide lower bound " + "on the largest per-sequence KV length in the batch and 0 means 'unknown'. It is a " + "provider-neutral performance hint; omitting it preserves the shape-(2) contract and disables " + "optimizations that require a lower bound unless the op reads exact lengths back from the device. " + "The op can otherwise obtain the upper bounds only by copying " "'cumulative_sequence_length' and 'past_seqlens' back from the device and synchronizing the " "stream on every call, which stalls the pipeline once per node per step and makes the op " "impossible to capture into a CUDA Graph. Schedulers already track these bounds on the host, so " "supplying them is normally free. When absent, the op falls back to the device readback. " - "The values are trusted: an under-sized bound violates the contract and may omit attention work.", + "The upper bounds are trusted: an under-sized bound violates the contract and may omit " + "attention work.", "S", OpSchema::Optional) .Output(0, diff --git a/onnxruntime/test/contrib_ops/paged_attention_op_test.cc b/onnxruntime/test/contrib_ops/paged_attention_op_test.cc index 05d62dcbeb5ed..791cd218f7792 100644 --- a/onnxruntime/test/contrib_ops/paged_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/paged_attention_op_test.cc @@ -20,9 +20,11 @@ #include "contrib_ops/cpu/bert/attention_common.h" #include "core/graph/model.h" #include "core/graph/node_attr_utils.h" +#include "core/providers/cuda/cuda_provider_options.h" #include "core/session/IOBinding.h" #include "core/session/inference_session.h" #include "default_providers.h" +#include "test/common/cuda_op_test_utils.h" #include "test/common/tensor_op_test_utils.h" #include "test/providers/provider_test_utils.h" #include "test/unittest_util/framework_test_utils.h" @@ -59,6 +61,25 @@ struct EndToEndCase { std::vector block_table; }; +struct IoBindingCase { + int batch_size = 1; + int num_heads = 1; + int kv_num_heads = 1; + int head_size = 8; + int block_size = 256; + int num_blocks = 2; + int max_num_blocks_per_seq = 1; + int past_seqlen = 4; + bool split_sensitive_values = false; + bool int8_cache = false; + bool enable_cuda_graph = false; + bool irregular_layout = false; + std::vector> replay_past_seqlens; + std::vector block_table; + std::vector attention_metadata; + std::string expected_error; +}; + // Softmax with causal masking: masked positions get -inf → 0 after exp. // Uses fp32 throughout to establish a reference the fp16 kernel is compared // against with a loose tolerance. @@ -247,17 +268,39 @@ void RunEndToEndCase(const EndToEndCase& c, std::unique_ptr void RunIoBindingCase(std::unique_ptr execution_provider, const char* provider_type, bool alias_cache_outputs, - bool omit_cache_outputs = false) { - constexpr int batch_size = 1; - constexpr int token_count = 1; - constexpr int num_heads = 1; - constexpr int kv_num_heads = 1; - constexpr int head_size = 8; - constexpr int block_size = 256; - constexpr int num_blocks = 2; - constexpr int past_seqlen = 4; - constexpr int hidden_size = num_heads * head_size; - constexpr int cache_elems = num_blocks * block_size * kv_num_heads * head_size; + bool omit_cache_outputs = false, + const IoBindingCase& c = IoBindingCase{}) { + const int batch_size = c.batch_size; + const int token_count = batch_size; + const int num_heads = c.num_heads; + const int kv_num_heads = c.kv_num_heads; + const int head_size = c.head_size; + const int block_size = c.block_size; + const int num_blocks = c.num_blocks; + const int max_num_blocks_per_seq = c.max_num_blocks_per_seq; + const int past_seqlen = c.past_seqlen; + const int hidden_size = num_heads * head_size; + const int kv_hidden_size = kv_num_heads * head_size; + const int cache_elems = num_blocks * block_size * kv_num_heads * head_size; + constexpr float cache_scale = 0.01f; + + ASSERT_FALSE(c.replay_past_seqlens.empty() && c.enable_cuda_graph); + for (const auto& replay_lengths : c.replay_past_seqlens) { + ASSERT_EQ(replay_lengths.size(), static_cast(batch_size)); + for (int32_t replay_length : replay_lengths) { + ASSERT_GE(replay_length, 0); + ASSERT_GT(max_num_blocks_per_seq, replay_length / block_size); + } + } + ASSERT_TRUE(!c.replay_past_seqlens.empty() || max_num_blocks_per_seq > past_seqlen / block_size); + ASSERT_LE(batch_size * max_num_blocks_per_seq, num_blocks); + ASSERT_EQ(num_heads % kv_num_heads, 0); + ASSERT_TRUE(c.block_table.empty() || + c.block_table.size() == static_cast(batch_size * max_num_blocks_per_seq)); + for (int32_t block_id : c.block_table) { + ASSERT_GE(block_id, 0); + ASSERT_LT(block_id, num_blocks); + } std::unordered_map domain_to_version = {{onnxruntime::kMSDomain, 1}}; std::vector model_specific_functions; @@ -283,33 +326,59 @@ void RunIoBindingCase(std::unique_ptr execution_provider, auto& query_arg = graph.GetOrCreateNodeArg("query", add_tensor_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT16, {token_count, hidden_size})); auto& key_arg = graph.GetOrCreateNodeArg("key", add_tensor_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT16, - {token_count, hidden_size})); + {token_count, kv_hidden_size})); auto& value_arg = graph.GetOrCreateNodeArg("value", add_tensor_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT16, - {token_count, hidden_size})); + {token_count, kv_hidden_size})); + const int cache_elem_type = c.int8_cache ? ONNX_NAMESPACE::TensorProto_DataType_INT8 + : ONNX_NAMESPACE::TensorProto_DataType_FLOAT16; auto& key_cache_arg = graph.GetOrCreateNodeArg( - "key_cache", add_tensor_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT16, + "key_cache", add_tensor_type(cache_elem_type, {num_blocks, block_size, kv_num_heads, head_size})); auto& value_cache_arg = graph.GetOrCreateNodeArg( - "value_cache", add_tensor_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT16, + "value_cache", add_tensor_type(cache_elem_type, {num_blocks, block_size, kv_num_heads, head_size})); auto& cumulative_sequence_length_arg = graph.GetOrCreateNodeArg( "cumulative_sequence_length", add_tensor_type(ONNX_NAMESPACE::TensorProto_DataType_INT32, {batch_size + 1})); auto& past_seqlens_arg = graph.GetOrCreateNodeArg( "past_seqlens", add_tensor_type(ONNX_NAMESPACE::TensorProto_DataType_INT32, {batch_size})); auto& block_table_arg = graph.GetOrCreateNodeArg( - "block_table", add_tensor_type(ONNX_NAMESPACE::TensorProto_DataType_INT32, {batch_size, 1})); + "block_table", + add_tensor_type(ONNX_NAMESPACE::TensorProto_DataType_INT32, {batch_size, max_num_blocks_per_seq})); auto& empty_optional_arg = graph.GetOrCreateNodeArg("", nullptr); + NodeArg* attention_metadata_arg = &empty_optional_arg; + if (!c.attention_metadata.empty()) { + attention_metadata_arg = &graph.GetOrCreateNodeArg( + "attention_metadata", + add_tensor_type(ONNX_NAMESPACE::TensorProto_DataType_INT32, + {static_cast(c.attention_metadata.size())})); + } + NodeArg* k_scale_arg = &empty_optional_arg; + NodeArg* v_scale_arg = &empty_optional_arg; + if (c.int8_cache) { + k_scale_arg = &graph.GetOrCreateNodeArg( + "k_scale", add_tensor_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT, {1})); + v_scale_arg = &graph.GetOrCreateNodeArg( + "v_scale", add_tensor_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT, {1})); + } std::vector input_defs = {&query_arg, &key_arg, &value_arg, &key_cache_arg, &value_cache_arg, &cumulative_sequence_length_arg, &past_seqlens_arg, &block_table_arg, - &empty_optional_arg, &empty_optional_arg}; + /*cos_cache=*/&empty_optional_arg, + /*sin_cache=*/&empty_optional_arg, + /*slot_mapping=*/&empty_optional_arg, + /*head_sink=*/&empty_optional_arg, + /*q_norm_weight=*/&empty_optional_arg, + /*k_norm_weight=*/&empty_optional_arg, + k_scale_arg, + v_scale_arg, + attention_metadata_arg}; auto& output_arg = graph.GetOrCreateNodeArg( "output", add_tensor_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT16, {token_count, hidden_size})); auto& key_cache_out_arg = graph.GetOrCreateNodeArg( - "key_cache_out", add_tensor_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT16, + "key_cache_out", add_tensor_type(cache_elem_type, {num_blocks, block_size, kv_num_heads, head_size})); auto& value_cache_out_arg = graph.GetOrCreateNodeArg( - "value_cache_out", add_tensor_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT16, + "value_cache_out", add_tensor_type(cache_elem_type, {num_blocks, block_size, kv_num_heads, head_size})); std::vector output_defs = {&output_arg}; if (!omit_cache_outputs) { @@ -323,6 +392,10 @@ void RunIoBindingCase(std::unique_ptr execution_provider, {"scale", utils::MakeAttribute("scale", 0.0f)}, {"do_rotary", utils::MakeAttribute("do_rotary", int64_t{0})}, }; + if (c.int8_cache) { + attrs.emplace("k_quant_type", utils::MakeAttribute("k_quant_type", std::string{"PER_TENSOR"})); + attrs.emplace("v_quant_type", utils::MakeAttribute("v_quant_type", std::string{"PER_TENSOR"})); + } auto& node = graph.AddNode("paged_attention", "PagedAttention", "IOBinding cache test", input_defs, output_defs, &attrs, onnxruntime::kMSDomain); node.SetExecutionProviderType(provider_type); @@ -355,6 +428,13 @@ void RunIoBindingCase(std::unique_ptr execution_provider, ASSERT_NE(device_alloc, nullptr); auto cpu_alloc = TestCPUExecutionProvider()->CreatePreferredAllocators()[0]; + OrtValue attention_metadata_value; + if (!c.attention_metadata.empty()) { + Tensor cpu_tensor(DataTypeImpl::GetType(), + TensorShape({static_cast(c.attention_metadata.size())}), + const_cast(c.attention_metadata.data()), cpu_alloc->Info()); + Tensor::InitOrtValue(std::move(cpu_tensor), attention_metadata_value); + } auto make_gpu_fp16 = [&](const std::vector& data, const TensorShape& shape) { Tensor cpu_tensor(DataTypeImpl::GetType(), shape, const_cast(data.data()), cpu_alloc->Info()); @@ -372,26 +452,141 @@ void RunIoBindingCase(std::unique_ptr execution_provider, Tensor::InitOrtValue(std::move(gpu_tensor), value); return value; }; + auto make_gpu_int8 = [&](const std::vector& data, const TensorShape& shape) { + Tensor cpu_tensor(DataTypeImpl::GetType(), shape, const_cast(data.data()), cpu_alloc->Info()); + Tensor gpu_tensor(DataTypeImpl::GetType(), shape, device_alloc); + ORT_THROW_IF_ERROR(execution_provider_ptr->GetDataTransfer()->CopyTensor(cpu_tensor, gpu_tensor)); + OrtValue value; + Tensor::InitOrtValue(std::move(gpu_tensor), value); + return value; + }; + auto make_gpu_float = [&](const std::vector& data, const TensorShape& shape) { + Tensor cpu_tensor(DataTypeImpl::GetType(), shape, const_cast(data.data()), cpu_alloc->Info()); + Tensor gpu_tensor(DataTypeImpl::GetType(), shape, device_alloc); + ORT_THROW_IF_ERROR(execution_provider_ptr->GetDataTransfer()->CopyTensor(cpu_tensor, gpu_tensor)); + OrtValue value; + Tensor::InitOrtValue(std::move(gpu_tensor), value); + return value; + }; + + std::vector block_table_data = c.block_table; + if (block_table_data.empty()) { + block_table_data.resize(batch_size * max_num_blocks_per_seq); + std::iota(block_table_data.begin(), block_table_data.end(), 0); + } std::vector query_data(token_count * hidden_size, MLFloat16(0.02f)); - std::vector key_data(token_count * hidden_size, MLFloat16(0.03f)); - std::vector value_data(token_count * hidden_size, MLFloat16(0.04f)); + std::vector key_data(token_count * kv_hidden_size, MLFloat16(0.03f)); + std::vector value_data(token_count * kv_hidden_size); + for (int b = 0; b < batch_size; ++b) { + const MLFloat16 value(0.04f + 0.02f * b); + std::fill_n(value_data.begin() + b * kv_hidden_size, kv_hidden_size, value); + } std::vector key_cache_data(cache_elems, MLFloat16(0.01f)); std::vector value_cache_data(cache_elems, MLFloat16(0.02f)); + if (c.irregular_layout) { + for (int b = 0; b < batch_size; ++b) { + for (int q_head = 0; q_head < num_heads; ++q_head) { + for (int dim = 0; dim < head_size; ++dim) { + const int index = (b * num_heads + q_head) * head_size + dim; + query_data[index] = MLFloat16(0.002f * static_cast((b * 5 + q_head * 3 + dim) % 11 - 5)); + } + } + for (int kv_head = 0; kv_head < kv_num_heads; ++kv_head) { + for (int dim = 0; dim < head_size; ++dim) { + const int index = (b * kv_num_heads + kv_head) * head_size + dim; + key_data[index] = + MLFloat16(0.003f * static_cast((b * 7 + kv_head * 5 + dim) % 13 - 6)); + value_data[index] = + MLFloat16(0.004f * static_cast((b * 3 + kv_head * 7 + dim) % 17 - 8)); + } + } + } + for (int block_id = 0; block_id < num_blocks; ++block_id) { + for (int slot = 0; slot < block_size; ++slot) { + for (int kv_head = 0; kv_head < kv_num_heads; ++kv_head) { + for (int dim = 0; dim < head_size; ++dim) { + const int index = CacheIndex(block_id, slot, kv_head, dim, + block_size, kv_num_heads, head_size); + key_cache_data[index] = MLFloat16( + 0.001f * static_cast((block_id * 3 + slot + kv_head * 5 + dim) % 13 - 6)); + value_cache_data[index] = MLFloat16( + 0.003f * static_cast((block_id * 5 + slot * 3 + kv_head * 7 + dim) % 17 - 8)); + } + } + } + } + } else if (c.split_sensitive_values) { + const int sequence_capacity = max_num_blocks_per_seq * block_size; + for (int b = 0; b < batch_size; ++b) { + const float low_value = -0.1f + 0.2f * b; + const float high_value = 0.1f + 0.2f * b; + for (int slot = 0; slot < sequence_capacity; ++slot) { + const MLFloat16 value(slot < past_seqlen / 2 ? low_value : high_value); + const int block_id = block_table_data[b * max_num_blocks_per_seq + slot / block_size]; + const int slot_offset = CacheIndex(block_id, slot % block_size, 0, 0, + block_size, kv_num_heads, head_size); + std::fill_n(value_cache_data.begin() + slot_offset, kv_num_heads * head_size, value); + } + } + } auto query_value = make_gpu_fp16(query_data, TensorShape({token_count, hidden_size})); - auto key_value = make_gpu_fp16(key_data, TensorShape({token_count, hidden_size})); - auto value_value = make_gpu_fp16(value_data, TensorShape({token_count, hidden_size})); - auto key_cache_value = make_gpu_fp16(key_cache_data, TensorShape({num_blocks, block_size, kv_num_heads, head_size})); - auto value_cache_value = make_gpu_fp16(value_cache_data, TensorShape({num_blocks, block_size, kv_num_heads, head_size})); - auto cumulative_sequence_length_value = make_gpu_int32({0, token_count}, TensorShape({batch_size + 1})); - auto past_seqlens_value = make_gpu_int32({past_seqlen}, TensorShape({batch_size})); - auto block_table_value = make_gpu_int32({0}, TensorShape({batch_size, 1})); + auto key_value = make_gpu_fp16(key_data, TensorShape({token_count, kv_hidden_size})); + auto value_value = make_gpu_fp16(value_data, TensorShape({token_count, kv_hidden_size})); + std::vector key_cache_int8; + std::vector value_cache_int8; + OrtValue key_cache_value; + OrtValue value_cache_value; + if (c.int8_cache) { + key_cache_int8.reserve(key_cache_data.size()); + value_cache_int8.reserve(value_cache_data.size()); + for (const auto value : key_cache_data) { + key_cache_int8.push_back(static_cast(std::round(value.ToFloat() / cache_scale))); + } + for (const auto value : value_cache_data) { + value_cache_int8.push_back(static_cast(std::round(value.ToFloat() / cache_scale))); + } + key_cache_value = make_gpu_int8(key_cache_int8, + TensorShape({num_blocks, block_size, kv_num_heads, head_size})); + value_cache_value = make_gpu_int8(value_cache_int8, + TensorShape({num_blocks, block_size, kv_num_heads, head_size})); + } else { + key_cache_value = make_gpu_fp16(key_cache_data, + TensorShape({num_blocks, block_size, kv_num_heads, head_size})); + value_cache_value = make_gpu_fp16(value_cache_data, + TensorShape({num_blocks, block_size, kv_num_heads, head_size})); + } + std::vector cumulative_sequence_length_data(batch_size + 1); + std::iota(cumulative_sequence_length_data.begin(), cumulative_sequence_length_data.end(), 0); + auto cumulative_sequence_length_value = + make_gpu_int32(cumulative_sequence_length_data, TensorShape({batch_size + 1})); + std::vector past_seqlens_data = + c.replay_past_seqlens.empty() ? std::vector(batch_size, past_seqlen) + : c.replay_past_seqlens.front(); + auto past_seqlens_value = make_gpu_int32(past_seqlens_data, TensorShape({batch_size})); + auto block_table_value = + make_gpu_int32(block_table_data, TensorShape({batch_size, max_num_blocks_per_seq})); auto output_value = make_gpu_fp16(std::vector(token_count * hidden_size), TensorShape({token_count, hidden_size})); - auto key_cache_out_value = make_gpu_fp16(key_cache_data, - TensorShape({num_blocks, block_size, kv_num_heads, head_size})); - auto value_cache_out_value = make_gpu_fp16(value_cache_data, - TensorShape({num_blocks, block_size, kv_num_heads, head_size})); + OrtValue key_cache_out_value; + OrtValue value_cache_out_value; + if (c.int8_cache) { + key_cache_out_value = make_gpu_int8(key_cache_int8, + TensorShape({num_blocks, block_size, kv_num_heads, head_size})); + value_cache_out_value = make_gpu_int8(value_cache_int8, + TensorShape({num_blocks, block_size, kv_num_heads, head_size})); + } else { + key_cache_out_value = make_gpu_fp16(key_cache_data, + TensorShape({num_blocks, block_size, kv_num_heads, head_size})); + value_cache_out_value = make_gpu_fp16(value_cache_data, + TensorShape({num_blocks, block_size, kv_num_heads, head_size})); + } + OrtValue k_scale_value; + OrtValue v_scale_value; + if (c.int8_cache) { + k_scale_value = make_gpu_float({cache_scale}, TensorShape({1})); + v_scale_value = make_gpu_float({cache_scale}, TensorShape({1})); + } std::unique_ptr io_binding; ASSERT_STATUS_OK(session.NewIOBinding(&io_binding)); @@ -403,18 +598,108 @@ void RunIoBindingCase(std::unique_ptr execution_provider, ASSERT_STATUS_OK(io_binding->BindInput("cumulative_sequence_length", cumulative_sequence_length_value)); ASSERT_STATUS_OK(io_binding->BindInput("past_seqlens", past_seqlens_value)); ASSERT_STATUS_OK(io_binding->BindInput("block_table", block_table_value)); + if (c.int8_cache) { + ASSERT_STATUS_OK(io_binding->BindInput("k_scale", k_scale_value)); + ASSERT_STATUS_OK(io_binding->BindInput("v_scale", v_scale_value)); + } + if (!c.attention_metadata.empty()) { + ASSERT_STATUS_OK(io_binding->BindInput("attention_metadata", attention_metadata_value)); + } ASSERT_STATUS_OK(io_binding->BindOutput("output", output_value)); if (!omit_cache_outputs) { ASSERT_STATUS_OK(io_binding->BindOutput("key_cache_out", alias_cache_outputs ? key_cache_value : key_cache_out_value)); ASSERT_STATUS_OK(io_binding->BindOutput("value_cache_out", alias_cache_outputs ? value_cache_value : value_cache_out_value)); } + const float scale = 1.0f / std::sqrt(static_cast(head_size)); + const size_t run_count = c.replay_past_seqlens.empty() ? 1 : c.replay_past_seqlens.size(); RunOptions run_options; - ASSERT_STATUS_OK(session.Run(run_options, *io_binding)); + if (c.enable_cuda_graph) { + ASSERT_STATUS_OK(run_options.config_options.AddConfigEntry("gpu_graph_id", "1")); + } + for (size_t run_index = 0; run_index < run_count; ++run_index) { + if (!c.replay_past_seqlens.empty()) { + past_seqlens_data = c.replay_past_seqlens[run_index]; + Tensor cpu_past_seqlens(DataTypeImpl::GetType(), TensorShape({batch_size}), + past_seqlens_data.data(), cpu_alloc->Info()); + ORT_THROW_IF_ERROR( + execution_provider_ptr->GetDataTransfer()->CopyTensor(cpu_past_seqlens, + *past_seqlens_value.GetMutable())); + Tensor cpu_cumulative(DataTypeImpl::GetType(), TensorShape({batch_size + 1}), + cumulative_sequence_length_data.data(), cpu_alloc->Info()); + ORT_THROW_IF_ERROR(execution_provider_ptr->GetDataTransfer()->CopyTensor( + cpu_cumulative, *cumulative_sequence_length_value.GetMutable())); + } + + const Status run_status = session.Run(run_options, *io_binding); + if (!c.expected_error.empty()) { + EXPECT_FALSE(run_status.IsOK()); + EXPECT_NE(run_status.ErrorMessage().find(c.expected_error), std::string::npos) + << run_status.ErrorMessage(); + return; + } + ASSERT_STATUS_OK(run_status); + + Tensor cpu_output(DataTypeImpl::GetType(), TensorShape({token_count, hidden_size}), cpu_alloc); + ORT_THROW_IF_ERROR(execution_provider_ptr->GetDataTransfer()->CopyTensor(output_value.Get(), cpu_output)); + for (int b = 0; b < batch_size; ++b) { + const int new_slot = past_seqlens_data[b]; + const int new_block_id = + block_table_data[b * max_num_blocks_per_seq + new_slot / block_size]; + for (int kv_head = 0; kv_head < kv_num_heads; ++kv_head) { + for (int dim = 0; dim < head_size; ++dim) { + const int cache_index = CacheIndex(new_block_id, new_slot % block_size, kv_head, dim, + block_size, kv_num_heads, head_size); + const int input_index = (b * kv_num_heads + kv_head) * head_size + dim; + key_cache_data[cache_index] = key_data[input_index]; + value_cache_data[cache_index] = value_data[input_index]; + } + } - Tensor cpu_output(DataTypeImpl::GetType(), TensorShape({token_count, hidden_size}), cpu_alloc); - ORT_THROW_IF_ERROR(execution_provider_ptr->GetDataTransfer()->CopyTensor(output_value.Get(), cpu_output)); - EXPECT_NE(cpu_output.Data()[0].ToFloat(), 0.0f); + const int gqa_factor = num_heads / kv_num_heads; + for (int q_head = 0; q_head < num_heads; ++q_head) { + const int kv_head = q_head / gqa_factor; + std::vector scores(new_slot + 1); + float max_score = -std::numeric_limits::infinity(); + for (int slot = 0; slot <= new_slot; ++slot) { + const int block_id = + block_table_data[b * max_num_blocks_per_seq + slot / block_size]; + float dot = 0.0f; + for (int dim = 0; dim < head_size; ++dim) { + const int query_index = (b * num_heads + q_head) * head_size + dim; + const int cache_index = CacheIndex(block_id, slot % block_size, kv_head, dim, + block_size, kv_num_heads, head_size); + dot += query_data[query_index].ToFloat() * key_cache_data[cache_index].ToFloat(); + } + scores[slot] = dot * scale; + max_score = std::max(max_score, scores[slot]); + } + float denominator = 0.0f; + for (float& score : scores) { + score = std::exp(score - max_score); + denominator += score; + } + for (int dim = 0; dim < head_size; ++dim) { + float numerator = 0.0f; + for (int slot = 0; slot <= new_slot; ++slot) { + const int block_id = + block_table_data[b * max_num_blocks_per_seq + slot / block_size]; + const int cache_index = CacheIndex(block_id, slot % block_size, kv_head, dim, + block_size, kv_num_heads, head_size); + numerator += scores[slot] * value_cache_data[cache_index].ToFloat(); + } + const int output_index = (b * num_heads + q_head) * head_size + dim; + EXPECT_NEAR(cpu_output.Data()[output_index].ToFloat(), + numerator / denominator, 2e-3f) + << "run=" << run_index << ", batch=" << b + << ", q_head=" << q_head << ", dim=" << dim; + } + } + } + } + if (c.enable_cuda_graph) { + EXPECT_TRUE(execution_provider_ptr->IsGraphCaptured(1)); + } const auto& outputs = io_binding->GetOutputs(); if (omit_cache_outputs) { @@ -424,20 +709,27 @@ void RunIoBindingCase(std::unique_ptr execution_provider, TensorShape({num_blocks, block_size, kv_num_heads, head_size}), cpu_alloc); ORT_THROW_IF_ERROR(execution_provider_ptr->GetDataTransfer()->CopyTensor(key_cache_value.Get(), cpu_key_cache)); ORT_THROW_IF_ERROR(execution_provider_ptr->GetDataTransfer()->CopyTensor(value_cache_value.Get(), cpu_value_cache)); - const size_t cache_update_offset = static_cast(past_seqlen * head_size); - EXPECT_NEAR(cpu_key_cache.Data()[cache_update_offset].ToFloat(), 0.03f, 1e-3f); - EXPECT_NEAR(cpu_value_cache.Data()[cache_update_offset].ToFloat(), 0.04f, 1e-3f); + for (int b = 0; b < batch_size; ++b) { + const int block_id = + block_table_data[b * max_num_blocks_per_seq + past_seqlen / block_size]; + const size_t cache_update_offset = + static_cast(CacheIndex(block_id, past_seqlen % block_size, 0, 0, + block_size, kv_num_heads, head_size)); + EXPECT_NEAR(cpu_key_cache.Data()[cache_update_offset].ToFloat(), 0.03f, 1e-3f); + EXPECT_NEAR(cpu_value_cache.Data()[cache_update_offset].ToFloat(), + 0.04f + 0.02f * b, 1e-3f); + } ASSERT_EQ(outputs.size(), 1u); return; } ASSERT_EQ(outputs.size(), 3u); if (alias_cache_outputs) { - EXPECT_EQ(outputs[1].Get().Data(), key_cache_value.Get().Data()); - EXPECT_EQ(outputs[2].Get().Data(), value_cache_value.Get().Data()); + EXPECT_EQ(outputs[1].Get().DataRaw(), key_cache_value.Get().DataRaw()); + EXPECT_EQ(outputs[2].Get().DataRaw(), value_cache_value.Get().DataRaw()); } else { - EXPECT_NE(outputs[1].Get().Data(), key_cache_value.Get().Data()); - EXPECT_NE(outputs[2].Get().Data(), value_cache_value.Get().Data()); + EXPECT_NE(outputs[1].Get().DataRaw(), key_cache_value.Get().DataRaw()); + EXPECT_NE(outputs[2].Get().DataRaw(), value_cache_value.Get().DataRaw()); } // Verify K/V scatter actually landed at slot `past_seqlen` in both caches. @@ -448,15 +740,28 @@ void RunIoBindingCase(std::unique_ptr execution_provider, // from "scatter silently didn't run". Downloading from the bound output // tensors covers both the aliased path (output backed by the input cache // buffer) and the non-aliased path (output backed by a separate buffer). - Tensor cpu_key_cache_out(DataTypeImpl::GetType(), + Tensor cpu_key_cache_out(c.int8_cache ? DataTypeImpl::GetType() : DataTypeImpl::GetType(), TensorShape({num_blocks, block_size, kv_num_heads, head_size}), cpu_alloc); - Tensor cpu_value_cache_out(DataTypeImpl::GetType(), + Tensor cpu_value_cache_out(c.int8_cache ? DataTypeImpl::GetType() : DataTypeImpl::GetType(), TensorShape({num_blocks, block_size, kv_num_heads, head_size}), cpu_alloc); ORT_THROW_IF_ERROR(execution_provider_ptr->GetDataTransfer()->CopyTensor(outputs[1].Get(), cpu_key_cache_out)); ORT_THROW_IF_ERROR(execution_provider_ptr->GetDataTransfer()->CopyTensor(outputs[2].Get(), cpu_value_cache_out)); - const size_t cache_update_offset = static_cast(past_seqlen * head_size); - EXPECT_NEAR(cpu_key_cache_out.Data()[cache_update_offset].ToFloat(), 0.03f, 1e-3f); - EXPECT_NEAR(cpu_value_cache_out.Data()[cache_update_offset].ToFloat(), 0.04f, 1e-3f); + for (int b = 0; b < batch_size; ++b) { + const int last_past_seqlen = past_seqlens_data[b]; + const int block_id = + block_table_data[b * max_num_blocks_per_seq + last_past_seqlen / block_size]; + const size_t cache_update_offset = + static_cast(CacheIndex(block_id, last_past_seqlen % block_size, 0, 0, + block_size, kv_num_heads, head_size)); + const float cached_key = c.int8_cache + ? cpu_key_cache_out.Data()[cache_update_offset] * cache_scale + : cpu_key_cache_out.Data()[cache_update_offset].ToFloat(); + const float cached_value = c.int8_cache + ? cpu_value_cache_out.Data()[cache_update_offset] * cache_scale + : cpu_value_cache_out.Data()[cache_update_offset].ToFloat(); + EXPECT_NEAR(cached_key, key_data[b * kv_hidden_size].ToFloat(), 1e-3f); + EXPECT_NEAR(cached_value, value_data[b * kv_hidden_size].ToFloat(), 1e-3f); + } } void RunEndToEndCaseOnAvailableProviders(const EndToEndCase& c) { @@ -487,7 +792,7 @@ TEST(PagedAttention, Cuda_AliasedCache_IOBinding) { RunIoBindingCase(DefaultCudaExecutionProvider(), kCudaExecutionProvider, true); } -TEST(PagedAttention, Cuda_DebugInfoIncludesDispatchBounds) { +TEST(PagedAttention, Cuda_AttentionMetadataShape2CompatibilityAndDispatchBounds) { ScopedEnvironmentVariables scoped_env_vars{ EnvVarMap{ {onnxruntime::contrib::attention::kDisableFlashAttention, "1"}, @@ -499,8 +804,11 @@ TEST(PagedAttention, Cuda_DebugInfoIncludesDispatchBounds) { GTEST_SKIP() << "CUDA EP not available."; } + IoBindingCase c; + c.attention_metadata = {1, 256}; + testing::internal::CaptureStdout(); - RunIoBindingCase(DefaultCudaExecutionProvider(), kCudaExecutionProvider, true); + RunIoBindingCase(DefaultCudaExecutionProvider(), kCudaExecutionProvider, true, false, c); const std::string debug_output = testing::internal::GetCapturedStdout(); EXPECT_NE(debug_output.find("Operator=PagedAttention"), std::string::npos) << debug_output; @@ -510,6 +818,208 @@ TEST(PagedAttention, Cuda_DebugInfoIncludesDispatchBounds) { EXPECT_NE(debug_output.find("EffectiveKvLengthBound=256"), std::string::npos) << debug_output; } +TEST(PagedAttention, Cuda_AttentionMetadataValidation) { + if (DefaultCudaExecutionProvider() == nullptr) { + GTEST_SKIP() << "CUDA EP not available."; + } + + struct InvalidMetadataCase { + std::vector metadata; + const char* expected_error; + }; + const std::vector cases = { + {{1, 256, -1}, "entries must be non-negative"}, + {{1, 128, 129}, "must not exceed max_kv_len_bound"}, + {{1, 0, 257}, "must not exceed max_kv_len_bound"}, + {{1, 256, 128, 64}, "must have shape (2) or (3)"}, + }; + + for (const auto& test_case : cases) { + IoBindingCase c; + c.attention_metadata = test_case.metadata; + c.expected_error = test_case.expected_error; + RunIoBindingCase(DefaultCudaExecutionProvider(), kCudaExecutionProvider, true, false, c); + } +} + +TEST(PagedAttention, Cuda_FlashSplitKvLongContext) { +#if defined(USE_FLASH_ATTENTION) + ScopedEnvironmentVariables scoped_env_vars{ + EnvVarMap{ + {onnxruntime::contrib::attention::kDisableFlashAttention, "0"}, + {onnxruntime::contrib::attention::kDisableMemoryEfficientAttention, "1"}, + {onnxruntime::contrib::attention::kDisableDecoderAttention, "1"}, + {onnxruntime::contrib::attention::kEnableAttentionKernelDebugInfo, "1"}}}; + + if (DefaultCudaExecutionProvider() == nullptr) { + GTEST_SKIP() << "CUDA EP not available."; + } + if (GetCudaArchitecture() < 800) { + GTEST_SKIP() << "Flash Attention requires compute capability 8.0 or later."; + } + + IoBindingCase c; + c.batch_size = 2; + c.num_heads = 4; + c.kv_num_heads = 2; + c.head_size = 128; + c.num_blocks = 32; + c.max_num_blocks_per_seq = 16; + c.irregular_layout = true; + c.replay_past_seqlens = {{767, 2047}}; + c.block_table.resize(c.batch_size * c.max_num_blocks_per_seq); + for (int i = 0; i < static_cast(c.block_table.size()); ++i) { + c.block_table[i] = (i * 13 + 7) % c.num_blocks; + } + c.attention_metadata = {1, 4096, 2048}; + + testing::internal::CaptureStdout(); + RunIoBindingCase(DefaultCudaExecutionProvider(), kCudaExecutionProvider, true, false, c); + const std::string debug_output = testing::internal::GetCapturedStdout(); + + EXPECT_NE(debug_output.find("SdpaKernel=FLASH_ATTENTION"), std::string::npos) << debug_output; + EXPECT_NE(debug_output.find("EffectiveKvLengthBound=4096"), std::string::npos) << debug_output; + const std::string split_prefix = "NumSplits="; + const size_t split_pos = debug_output.find(split_prefix); + ASSERT_NE(split_pos, std::string::npos) << debug_output; + EXPECT_GT(std::stoi(debug_output.substr(split_pos + split_prefix.size())), 1) << debug_output; +#else + GTEST_SKIP() << "Flash Attention is not enabled in this build."; +#endif +} + +TEST(PagedAttention, Cuda_FlashSplitKvCudaGraphReplay) { +#if defined(USE_FLASH_ATTENTION) + ScopedEnvironmentVariables scoped_env_vars{ + EnvVarMap{ + {onnxruntime::contrib::attention::kDisableFlashAttention, "0"}, + {onnxruntime::contrib::attention::kDisableMemoryEfficientAttention, "1"}, + {onnxruntime::contrib::attention::kDisableDecoderAttention, "1"}, + {onnxruntime::contrib::attention::kEnableAttentionKernelDebugInfo, "1"}}}; + + if (DefaultCudaExecutionProvider() == nullptr) { + GTEST_SKIP() << "CUDA EP not available."; + } + if (GetCudaArchitecture() < 800) { + GTEST_SKIP() << "Flash Attention requires compute capability 8.0 or later."; + } + + OrtCUDAProviderOptionsV2 provider_options{}; + provider_options.do_copy_in_default_stream = true; + provider_options.use_tf32 = false; + provider_options.enable_cuda_graph = true; + + IoBindingCase c; + c.batch_size = 2; + c.num_heads = 2; + c.kv_num_heads = 1; + c.head_size = 128; + c.num_blocks = 256; + c.max_num_blocks_per_seq = 128; + c.irregular_layout = true; + c.enable_cuda_graph = true; + c.replay_past_seqlens = { + {512, 512}, + {513, 1024}, + {1024, 4096}, + {2048, 8192}, + }; + c.attention_metadata = {1, 32768, 513}; + + testing::internal::CaptureStdout(); + RunIoBindingCase(CudaExecutionProviderWithOptions(&provider_options), + kCudaExecutionProvider, true, false, c); + const std::string debug_output = testing::internal::GetCapturedStdout(); + + EXPECT_NE(debug_output.find("SdpaKernel=FLASH_ATTENTION"), std::string::npos) << debug_output; + const std::string split_prefix = "NumSplits="; + const size_t split_pos = debug_output.find(split_prefix); + ASSERT_NE(split_pos, std::string::npos) << debug_output; + EXPECT_GT(std::stoi(debug_output.substr(split_pos + split_prefix.size())), 1) << debug_output; +#else + GTEST_SKIP() << "Flash Attention is not enabled in this build."; +#endif +} + +TEST(PagedAttention, Cuda_FlashSplitKvInt8Cache) { +#if defined(USE_FLASH_ATTENTION) + ScopedEnvironmentVariables scoped_env_vars{ + EnvVarMap{ + {onnxruntime::contrib::attention::kDisableFlashAttention, "0"}, + {onnxruntime::contrib::attention::kDisableMemoryEfficientAttention, "1"}, + {onnxruntime::contrib::attention::kDisableDecoderAttention, "1"}, + {onnxruntime::contrib::attention::kEnableAttentionKernelDebugInfo, "1"}}}; + + if (DefaultCudaExecutionProvider() == nullptr) { + GTEST_SKIP() << "CUDA EP not available."; + } + if (GetCudaArchitecture() < 800) { + GTEST_SKIP() << "Flash Attention requires compute capability 8.0 or later."; + } + + IoBindingCase c; + c.batch_size = 2; + c.num_heads = 2; + c.kv_num_heads = 1; + c.head_size = 128; + c.num_blocks = 32; + c.max_num_blocks_per_seq = 16; + c.past_seqlen = 2047; + c.split_sensitive_values = true; + c.int8_cache = true; + c.attention_metadata = {1, 2048, 2048}; + + testing::internal::CaptureStdout(); + RunIoBindingCase(DefaultCudaExecutionProvider(), kCudaExecutionProvider, true, false, c); + const std::string debug_output = testing::internal::GetCapturedStdout(); + + EXPECT_NE(debug_output.find("SdpaKernel=FLASH_ATTENTION"), std::string::npos) << debug_output; + const std::string split_prefix = "NumSplits="; + const size_t split_pos = debug_output.find(split_prefix); + ASSERT_NE(split_pos, std::string::npos) << debug_output; + EXPECT_GT(std::stoi(debug_output.substr(split_pos + split_prefix.size())), 1) << debug_output; +#else + GTEST_SKIP() << "Flash Attention is not enabled in this build."; +#endif +} + +TEST(PagedAttention, Cuda_FlashSplitKvSkipsShortReplayRange) { +#if defined(USE_FLASH_ATTENTION) + ScopedEnvironmentVariables scoped_env_vars{ + EnvVarMap{ + {onnxruntime::contrib::attention::kDisableFlashAttention, "0"}, + {onnxruntime::contrib::attention::kDisableMemoryEfficientAttention, "1"}, + {onnxruntime::contrib::attention::kDisableDecoderAttention, "1"}, + {onnxruntime::contrib::attention::kEnableAttentionKernelDebugInfo, "1"}}}; + + if (DefaultCudaExecutionProvider() == nullptr) { + GTEST_SKIP() << "CUDA EP not available."; + } + if (GetCudaArchitecture() < 800) { + GTEST_SKIP() << "Flash Attention requires compute capability 8.0 or later."; + } + + IoBindingCase c; + c.num_heads = 2; + c.kv_num_heads = 1; + c.head_size = 128; + c.num_blocks = 16; + c.max_num_blocks_per_seq = 16; + c.past_seqlen = 127; + c.attention_metadata = {1, 2048, 128}; + + testing::internal::CaptureStdout(); + RunIoBindingCase(DefaultCudaExecutionProvider(), kCudaExecutionProvider, true, false, c); + const std::string debug_output = testing::internal::GetCapturedStdout(); + + EXPECT_NE(debug_output.find("SdpaKernel=FLASH_ATTENTION"), std::string::npos) << debug_output; + EXPECT_NE(debug_output.find("EffectiveKvLengthBound=2048"), std::string::npos) << debug_output; + EXPECT_NE(debug_output.find("NumSplits=1"), std::string::npos) << debug_output; +#else + GTEST_SKIP() << "Flash Attention is not enabled in this build."; +#endif +} + TEST(PagedAttention, WebGpu_AliasedCache_IOBinding) { if (DefaultWebGpuExecutionProvider() == nullptr) { GTEST_SKIP() << "WebGPU EP not available."; diff --git a/onnxruntime/test/providers/cuda/test_cases/attention_split_heuristic_test.cc b/onnxruntime/test/providers/cuda/test_cases/attention_split_heuristic_test.cc index 78293c3d2bbad..a289010a957e3 100644 --- a/onnxruntime/test/providers/cuda/test_cases/attention_split_heuristic_test.cc +++ b/onnxruntime/test/providers/cuda/test_cases/attention_split_heuristic_test.cc @@ -68,6 +68,44 @@ TEST(FlashAttentionTest, GetNumSplitsHandlesZeroKeyTiles) { #endif } +TEST(FlashAttentionTest, GetNumSplitsUsesLongContextParallelism) { +#if defined(USE_FLASH_ATTENTION) + const auto [num_splits, softmax_lse_accum_bytes, out_accum_bytes] = + flash::get_num_splits_and_buffer_sizes( + 1, // batch_size + 1, // seqlen_q + 4096, // seqlen_k + 2, // num_heads + 64, // head_size + 108); // num_SMs + + EXPECT_EQ(num_splits, 16U); + EXPECT_GT(softmax_lse_accum_bytes, 0U); + EXPECT_GT(out_accum_bytes, 0U); +#else + GTEST_SKIP() << "Flash Attention is not enabled in this build."; +#endif +} + +TEST(FlashAttentionTest, GetNumSplitsAvoidsShortContextOverhead) { +#if defined(USE_FLASH_ATTENTION) + const auto [num_splits, softmax_lse_accum_bytes, out_accum_bytes] = + flash::get_num_splits_and_buffer_sizes( + 1, // batch_size + 1, // seqlen_q + 128, // KV sequence length + 2, // num_heads + 64, // head_size + 108); // num_SMs + + EXPECT_EQ(num_splits, 0U); + EXPECT_EQ(softmax_lse_accum_bytes, 0U); + EXPECT_EQ(out_accum_bytes, 0U); +#else + GTEST_SKIP() << "Flash Attention is not enabled in this build."; +#endif +} + TEST(LeanAttentionTest, GetNumSplitsHandlesZeroSmCount) { #if defined(USE_LEAN_ATTENTION) const auto [num_splits, softmax_lse_accum_bytes, out_accum_bytes, sync_flag_bytes,