Sync with Microsoft ONNX Runtime - 19082026

docs/ContribOperators.md

@@ -4532,7 +4532,7 @@ This version of the operator has been available since version 1 of the 'com.micr
 <dt><tt>v_scale</tt> (optional) : T_KV_SCALE</dt>
 <dd>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).</dd>
 <dt><tt>attention_metadata</tt> (optional) : S</dt>
-<dd>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.</dd>
+<dd>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.</dd>
 </dl>
 
 #### Outputs (1 - 3)
@@ -7286,5 +7286,3 @@ No versioning maintained for experimental ops.
 <dt><tt>T</tt> : tensor(float)</dt>
 <dd>Constrain input and output types to float32 tensors.</dd>
 </dl>
-
-

docs/contrib_ops/cuda/matmul_block_scaled_fp4.md

@@ -28,6 +28,7 @@ Source files:
 3. [Dispatch Chain](#3-dispatch-chain)
 4. [Decode Path - Fused GEMV](#4-decode-path---fused-gemv)
 5. [Default Path - Dequantize + cuBLAS](#5-default-path---dequantize--cublas)
+   - [5.1 Vectorized dequantization](#51-vectorized-dequantization)
 6. [Native SM120 FP4 x FP4 Path](#6-native-sm120-fp4-x-fp4-path)
 7. [PrePack](#7-prepack)
 8. [Environment Variables](#8-environment-variables)
@@ -220,6 +221,31 @@ This path keeps full-precision activations and runs on CUDA devices with NVFP4
 conversion intrinsic support in the configured CUDA toolkit. It is the default
 prefill path when the SM120 native environment variable is not enabled.
 
+### 5.1 Vectorized dequantization
+
+When `K % 8 == 0` and `block_size` is even - the layout every real NVFP4 model
+uses - `LaunchDequantizeNvFp4` picks `DequantizeNvFp4Vec8Kernel` instead of the
+scalar kernel. Each thread owns exactly one 8-element K chunk of one row, so a
+warp issues one contiguous 128-byte packed load and one contiguous 512-byte
+store. Widening the per-thread chunk beyond one `uint4` store was measured to be
+about 2x slower because each store instruction then strides across lanes
+(1.9 vs 3.9 TB/s on H200). The row index comes from `blockIdx.y`, which removes
+the 64-bit division of the scalar kernel, `weight_scale_2` is hoisted into a
+register, and codes are decoded with the same branch-free `Fp4Cvt` `prmt`
+lookup the decode GEMV uses rather than the software-emulated
+`__nv_cvt_fp4x2_to_halfraw2()`.
+
+The output is bitwise identical to the scalar kernel. Measured on H200 for
+`M = 1024`, BF16, `block_size = 16` (median dequant kernel time):
+
+| N | K | scalar | vectorized | speedup |
+|---:|---:|---:|---:|---:|
+| 4096 | 4096 | 60.7 us | 15.5 us | 3.93x |
+| 6144 | 2048 | 46.1 us | 12.1 us | 3.81x |
+| 2048 | 6144 | 46.2 us | 11.9 us | 3.88x |
+
+The scalar kernel remains for odd `block_size` or `K % 8 != 0`.
+
 ---
 
 ## 6. Native SM120 FP4 x FP4 Path
@@ -304,6 +330,11 @@ CUDA_VISIBLE_DEVICES=0 "$ORT_BUILD/onnxruntime_provider_test" \
   --gtest_filter='MatMulBlockQuantizedFp4WeightOpTest.*'
 ```
 
+The `Gemv*` cases cover the decode path and the `PrefillDequant*` cases cover the
+dequantize + cuBLAS path, with `M > 8` so the GEMV is skipped: `*Vectorized*` for
+`DequantizeNvFp4Vec8Kernel` and `OddBlockSize` / `KNotMultipleOf8` for the scalar
+fallback, each in both FP16 and BF16.
+
 Python harness examples:
 
 ```bash

docs/contrib_ops/cuda/matmul_block_scaled_fp8.md

@@ -37,6 +37,7 @@ Source files:
 3. [Dispatch Chain](#3-dispatch-chain)
 4. [Decode Path - Fused GEMV](#4-decode-path---fused-gemv)
 5. [Default Path - Dequantize + cuBLAS](#5-default-path---dequantize--cublas)
+   - [5.1 Tiling the dequantization scratch over N](#51-tiling-the-dequantization-scratch-over-n)
 6. [Optional W8A8 Activation Path](#6-optional-w8a8-activation-path)
 7. [Testing and Benchmarking](#7-testing-and-benchmarking)
 
@@ -196,6 +197,32 @@ back `LaunchDequantizeBlockScaledFp8`:
 This keeps the GEMM in the activation type and runs on any CUDA architecture
 with FP8 conversion intrinsics (CUDA >= 11.8). It is the default prefill path.
 
+### 5.1 Tiling the dequantization scratch over N
+
+A full `[N, K]` scratch is 2.37 GiB for a 248320 x 5120 LM head, and it is
+allocated for the whole GEMM even though the GEMM reads it once. The scratch is
+therefore capped and the dequantize + GEMM pair is run over N tiles that fit
+inside the cap. Because the row-major `[M, N]` output is column-major `[N, M]`
+to cuBLAS, an N tile is a plain row offset into `Y`, so no extra copy is needed:
+
+```
+for n_offset in 0, tile_rows, 2 * tile_rows, ...:
+    dequantize B[n_offset : n_offset + rows, :] into the scratch
+    cublasGemmHelper(...) writing Y + n_offset
+```
+
+`ORT_FP8_DEQUANT_SCRATCH_MIB` sets the cap in MiB (default 256). The default was
+chosen by sweeping it on Qwen3.8-27B at an 8K prompt:
+
+| cap | peak memory (MiB) | reduction vs. untiled (MiB) | TTFT change |
+|---|---:|---:|---:|
+| untiled | 32609 | baseline | baseline |
+| 1 GiB | 29509 | -3100 | +5.2% |
+| **256 MiB** | **28503** | **-4106** | **+1.1%** |
+| 128 MiB | 28503 | -4106 | +13.9% |
+
+Shapes small enough to fit the cap take a single tile and are unaffected.
+
 ---
 
 ## 6. Optional W8A8 Activation Path

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

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());
     }
   }

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

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;

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,