From d8d0795b507ea2f7fd956846191483f44d00cec4 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 29 Oct 2026 00:09:56 +0000 Subject: [PATCH 01/14] [CUDA] GatedAdd: trigger the programmatic-dependent-launch successor early in ProjectedSigmoidGateMulKernel ProjectedSigmoidGateMulKernel feeds the hyper-connection down kernel, which is launched with programmatic dependent launch and only reads the gate-mul output after griddepcontrol.wait. Issue griddepcontrol.launch_dependents at kernel start (sm_90+) so the successor's prologue (constant weight prefetch) overlaps with this kernel. It is a no-op for plain launches and does not change any result (generated-token hashes identical in 45/45 Qwen3.8 decode runs). Together with the same trigger in the fpA_intB GEMV kernels (#33689) it shortens an S=8 decode cycle by 1.4-1.7%. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/contrib_ops/cuda/bert/gated_add_impl.cu | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.cu b/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.cu index 82aa1baf8d463..8f6139e456c30 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.cu @@ -56,6 +56,11 @@ __global__ void ProjectedSigmoidGateMulKernel( const __nv_bfloat16* attention, const __nv_bfloat16* activation, const __nv_bfloat16* weight, const __nv_bfloat16* residual, __nv_bfloat16* output, int64_t input_width, int64_t output_width) { + // Let a programmatic-dependent-launch successor (the hyper-connection down kernel) start and prefetch its constant + // weights while this kernel runs; it reads our output only after griddepcontrol.wait. No-op for plain launches. +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); +#endif const int64_t row = blockIdx.x; const float gate = ProjectedRowGate(activation, weight, row, input_width); const int64_t row_base = row * output_width; From a2de846c492ff9a6fabeaa3e0df311a568046791 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Wed, 28 Oct 2026 13:01:48 +0000 Subject: [PATCH 02/14] [CUDA] DynamicSparseAttention: tensor-core prefill kernel for the INT8 KV cache The selected_only prefill path of an INT8 cache (bfloat16 query / output, head_size 256) dequantized every gathered K / V tile to bf16 in shared memory (extra pass, extra smem round trip, extra barriers) before the bf16 MMAs. The new kernel keeps the gathered rows int8 in shared memory (cp.async, double buffered) and converts the ldmatrix fragments to fp16 exactly in registers (magic-number conversion). The per-(kv head, channel) key scale is folded into the query once per CTA (Q' = fp16(Q * key_scale)); the value scale is applied to the output once. Scores stay in registers, only the row maxima and the fp16 probability tile cross warps (3 barriers per 64-candidate tile). fp32 accumulation, GQA group <= 16 padded to the m16n8k16 M dimension. Dispatch: INT8 cache + prefill kernel + 16-byte aligned present K / V and scales + 8-byte aligned output + enough opt-in shared memory; ORT_DSA_PREFILL_INT8_TC=0 restores the previous kernel. Standalone harness (rows 8192, 2048 candidates, 2 kv heads, group 12): 6.1 ms -> 2.15 ms per launch (2.83x); rms error vs an fp64 reference 1.6e-3 (old kernel 1.7e-3, 2.7e-3 for peaky softmax). Qwen3.8 (12 full-attention layers) TTFT, one round on the three packages: -8.4..-9.5% at 8192 and -7.6..-9.4% at 32768 prompt tokens, decode unchanged. Tests: Int8CachePrefillTensorCoreKernel_CUDA (heads 2 / 12 / 24 / 32 x candidate counts 1 / 64 / 65 / 509 x invalid indices, new and previous kernel), 8-byte K / V and 4-byte output offsets in the misaligned-buffer cases. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../bert/dynamic_sparse_attention_impl.cu | 21 +- .../dynamic_sparse_attention_prefill_int8.cuh | 477 ++++++++++++++++++ .../dynamic_sparse_attention_op_test.cc | 52 +- 3 files changed, 544 insertions(+), 6 deletions(-) create mode 100644 onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu index a612c7071aab9..88d738f026a56 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu @@ -15,6 +15,7 @@ #include #include "contrib_ops/cuda/bert/dynamic_sparse_attention_prefill.cuh" +#include "contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh" #include "core/common/safeint.h" #include "core/providers/cuda/cuda_common.h" @@ -2073,6 +2074,7 @@ Status LaunchDynamicSparseAttention( grouped_decode_threads_supported && UseGroupedDecodeAttention(parameters, sizeof(T), max_shared_memory_per_block, data.quantized_cache); bool use_prefill_kernel = false; + bool use_int8_tc_prefill = false; if constexpr (std::is_same_v) { if (!use_grouped_decode) { const bool enabled = ParseEnvironmentVariableWithDefault("ORT_DSA_PREFILL_KERNEL", 1) != 0; @@ -2101,6 +2103,13 @@ Status LaunchDynamicSparseAttention( parameters.num_heads % parameters.kv_num_heads == 0 && parameters.num_heads / parameters.kv_num_heads <= dsa_prefill::kMaxGroup && static_cast(optin_shared) >= dsa_prefill::kSharedBytes; + // The tensor-core INT8 variant converts the cached rows in registers; it needs 16-byte cache rows and 8-byte + // output vectors (the BF16-dequantizing kernel above is the fallback). + use_int8_tc_prefill = + use_prefill_kernel && data.quantized_cache && + ParseEnvironmentVariableWithDefault("ORT_DSA_PREFILL_INT8_TC", 1) != 0 && + IsAligned(data.present_key, 16) && IsAligned(data.present_value, 16) && IsAligned(data.output, 8) && + static_cast(optin_shared) >= dsa_prefill::kInt8SharedBytes; } } if (use_grouped_decode) { @@ -2249,7 +2258,17 @@ Status LaunchDynamicSparseAttention( if constexpr (std::is_same_v) { const dim3 grid(static_cast(parameters.batch_size * parameters.sequence_length), static_cast(parameters.kv_num_heads)); - if (data.quantized_cache) { + if (use_int8_tc_prefill) { + CUDA_RETURN_IF_ERROR(cudaFuncSetAttribute( + dsa_prefill::DynamicSparseAttentionPrefillInt8Kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(dsa_prefill::kInt8SharedBytes))); + dsa_prefill::DynamicSparseAttentionPrefillInt8Kernel<<>>( + data.prepared_query, reinterpret_cast(data.present_key), + reinterpret_cast(data.present_value), data.selected_indices, data.selected_counts, + data.seqlens_k, data.output, parameters.sequence_length, parameters.num_heads, parameters.kv_num_heads, + parameters.cache_capacity, parameters.max_selected, parameters.scale, data.key_scale, data.value_scale); + } else if (data.quantized_cache) { CUDA_RETURN_IF_ERROR(cudaFuncSetAttribute( dsa_prefill::DynamicSparseAttentionPrefillKernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(dsa_prefill::kSharedBytes))); diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh new file mode 100644 index 0000000000000..dfbdd7ae8afc7 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh @@ -0,0 +1,477 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +// +// INT8-cache variant of the tensor-core prefill kernel of DynamicSparseAttention (selected_only, main KV source, +// bfloat16 query / output, head_size 256). One CTA handles one (query row, KV head); the heads of the GQA group form +// the M dimension of an m16n8k16 MMA (padded to 16). +// +// The scales of the INT8 cache are per (kv head, channel) constants, so they leave the tile loop: +// * the key scale is folded into the query once per CTA: Q' = Q * key_scale (rounded once to fp16), S = Q' K_int8; +// * the value scale is applied to the output once per CTA: O = (P V_int8) * value_scale. +// The gathered int8 rows stay int8 in shared memory (cp.async, double buffered, one tile ahead). A fragment read with +// ldmatrix is converted to fp16 exactly in registers (1024 + (byte ^ 0x80) - 1152, "magic number" conversion) and +// consumed by fp16 MMAs with fp32 accumulation, so the per-tile dequantization pass, its shared memory round trip and +// its barriers of the BF16-dequantizing kernel disappear. The scores stay in registers; the only cross-warp exchanges +// per tile are the row maxima and the fp16 probability tile. +// +// The channel (reduction) dimension of S = Q' K^T is permuted so that one 4-byte ldmatrix element of a K row supplies +// both B registers of an MMA k step: slot 2t+e <-> channel 16j + 4t + e and slot 2t+8+e <-> channel 16j + 4t + 2 + e, +// with the same permutation applied to the A fragments of Q'. Likewise ldmatrix.trans on the V tile delivers channels +// 2n and 2n + 1 of two consecutive candidates, i.e. the B registers of two n tiles whose columns are the even / odd +// channels. + +#pragma once + +#include "contrib_ops/cuda/bert/dynamic_sparse_attention_prefill.cuh" + +#include + +namespace onnxruntime { +namespace contrib { +namespace cuda { +namespace dsa_prefill { + +constexpr int kInt8RowStride = kHeadDim + 16; // bytes; 272 = 17 x 16 keeps ldmatrix conflict free +constexpr int kInt8TileBytes = kTile * kInt8RowStride; +constexpr int kInt8QStride = kHeadDim + 8; // fp16 elements of the staged query rows +constexpr int kInt8PStride = kTile + 8; // fp16 elements +constexpr int kIdxStages = 3; + +// sK8[2] + sV8[2] (sQ aliases sK8[0] before the tile loop) + candidate indices + row maxima / sums + probabilities. +constexpr size_t kInt8SharedBytes = 4 * static_cast(kInt8TileBytes) + kIdxStages * kTile * sizeof(int32_t) + + kMaxGroup * 4 * sizeof(float) + kMaxGroup * kInt8PStride * sizeof(__half); + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 +__device__ __forceinline__ void CpAsync16Zfill(uint32_t smem_dst, const void* gmem_src, bool valid) { + const int src_size = valid ? 16 : 0; + asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" ::"r"(smem_dst), "l"(gmem_src), "r"(src_size)); +} + +__device__ __forceinline__ void LdmatrixX4Addr(uint32_t (&r)[4], uint32_t addr) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(addr)); +} + +__device__ __forceinline__ void LdmatrixX4TransAddr(uint32_t (&r)[4], uint32_t addr) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0,%1,%2,%3}, [%4];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(addr)); +} + +__device__ __forceinline__ void MmaF16(float (&c)[4], const uint32_t (&a)[4], uint32_t b0, uint32_t b1) { + asm("mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n" + : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1)); +} + +__device__ __forceinline__ float FastExp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;\n" : "=f"(y) : "f"(x)); + return y; +} + +// Two bytes b0, b1 (offset binary u = b ^ 0x80 in [0, 255]) placed in the mantissa of fp16 1024: 1024 + u; subtracting +// 1152 gives the signed value exactly. +__device__ __forceinline__ uint32_t Half2MinusBias(uint32_t x) { + uint32_t r; + asm("sub.rn.f16x2 %0, %1, %2;\n" : "=r"(r) : "r"(x), "r"(0x64806480u)); + return r; +} + +// bytes (0, 1) -> lo, bytes (2, 3) -> hi (two fp16 each). +__device__ __forceinline__ void Int8x4ToHalf2Adjacent(uint32_t w, uint32_t& lo, uint32_t& hi) { + const uint32_t x = w ^ 0x80808080u; + uint32_t a; + uint32_t b; + asm("prmt.b32 %0, %1, %2, 0x5140;\n" : "=r"(a) : "r"(x), "r"(0x64646464u)); + asm("prmt.b32 %0, %1, %2, 0x5342;\n" : "=r"(b) : "r"(x), "r"(0x64646464u)); + lo = Half2MinusBias(a); + hi = Half2MinusBias(b); +} + +// bytes (0, 2) -> even, bytes (1, 3) -> odd (two fp16 each). +__device__ __forceinline__ void Int8x4ToHalf2Strided(uint32_t w, uint32_t& even, uint32_t& odd) { + const uint32_t x = w ^ 0x80808080u; + uint32_t a; + uint32_t b; + asm("prmt.b32 %0, %1, %2, 0x5240;\n" : "=r"(a) : "r"(x), "r"(0x64646464u)); + asm("prmt.b32 %0, %1, %2, 0x5341;\n" : "=r"(b) : "r"(x), "r"(0x64646464u)); + even = Half2MinusBias(a); + odd = Half2MinusBias(b); +} +#endif + +// grid = (batch * sequence_length, kv_num_heads), block = kThreads, dynamic smem = kInt8SharedBytes. +// query: [rows, num_heads, 256] (already normalized / rotated); key, value: int8 [batch, kv_heads, capacity, 256]; +// key_scale, value_scale: float [kv_heads, 256]. +__global__ void __launch_bounds__(kThreads, 3) DynamicSparseAttentionPrefillInt8Kernel( + const __nv_bfloat16* __restrict__ query, + const int8_t* __restrict__ key, + const int8_t* __restrict__ value, + const int32_t* __restrict__ selected_indices, + const int32_t* __restrict__ selected_counts, + const int32_t* __restrict__ seqlens_k, + __nv_bfloat16* __restrict__ output, + int sequence_length, + int num_heads, + int kv_num_heads, + int capacity, + int max_selected, + float scale, + const float* __restrict__ key_scale, + const float* __restrict__ value_scale) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + extern __shared__ __align__(16) unsigned char smem_raw[]; + unsigned char* sK8 = smem_raw; + unsigned char* sV8 = smem_raw + 2 * kInt8TileBytes; + int32_t* sIdx = reinterpret_cast(smem_raw + 4 * kInt8TileBytes); + float* sM = reinterpret_cast(sIdx + kIdxStages * kTile); + __half* sP = reinterpret_cast<__half*>(sM + kMaxGroup * 4); + __half* sQ = reinterpret_cast<__half*>(smem_raw); + + const int tid = threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + const int r0 = lane >> 2; // MMA row of the C / A fragments (and row + 8) + const int t = lane & 3; + const int row = blockIdx.x; + const int kv_head = blockIdx.y; + const int batch = row / sequence_length; + const int s = row - batch * sequence_length; + const int group = num_heads / kv_num_heads; + + const int64_t total_length = static_cast(seqlens_k[batch]) + 1; + const bool valid_length = total_length >= sequence_length && total_length <= capacity; + const int query_position = valid_length ? static_cast(total_length) - sequence_length + s : -1; + const int total = valid_length ? static_cast(total_length) : 0; + const int count = valid_length ? max(0, min(selected_counts[row], max_selected)) : 0; + const int32_t* row_indices = selected_indices + static_cast(row) * max_selected; + const int64_t kv_head_offset = (static_cast(batch) * kv_num_heads + kv_head) * capacity * kHeadDim; + const int8_t* key_base = key + kv_head_offset; + const int8_t* value_base = value + kv_head_offset; + + // Q' = Q * key_scale as fp16 (clamped so that it cannot overflow); the padding rows are zero so their (ignored) + // scores stay finite. + { + const int chunk = tid & 31; + const float4 k0 = *reinterpret_cast(key_scale + kv_head * kHeadDim + chunk * 8); + const float4 k1 = *reinterpret_cast(key_scale + kv_head * kHeadDim + chunk * 8 + 4); + const __nv_bfloat16* query_group = query + (static_cast(row) * num_heads + kv_head * group) * kHeadDim; +#pragma unroll + for (int i = 0; i < kMaxGroup / 4; ++i) { + const int r = (tid >> 5) + 4 * i; + uint4 packed = make_uint4(0, 0, 0, 0); + if (r < group) { + const uint4 q = *reinterpret_cast(query_group + static_cast(r) * kHeadDim + chunk * 8); + const float2 q01 = __bfloat1622float2(*reinterpret_cast(&q.x)); + const float2 q23 = __bfloat1622float2(*reinterpret_cast(&q.y)); + const float2 q45 = __bfloat1622float2(*reinterpret_cast(&q.z)); + const float2 q67 = __bfloat1622float2(*reinterpret_cast(&q.w)); + auto clamp = [](float x) { return fminf(fmaxf(x, -65504.0f), 65504.0f); }; + const __half2 h01 = __floats2half2_rn(clamp(q01.x * k0.x), clamp(q01.y * k0.y)); + const __half2 h23 = __floats2half2_rn(clamp(q23.x * k0.z), clamp(q23.y * k0.w)); + const __half2 h45 = __floats2half2_rn(clamp(q45.x * k1.x), clamp(q45.y * k1.y)); + const __half2 h67 = __floats2half2_rn(clamp(q67.x * k1.z), clamp(q67.y * k1.w)); + packed.x = *reinterpret_cast(&h01); + packed.y = *reinterpret_cast(&h23); + packed.z = *reinterpret_cast(&h45); + packed.w = *reinterpret_cast(&h67); + } + *reinterpret_cast(sQ + r * kInt8QStride + chunk * 8) = packed; + } + } + __syncthreads(); + + // A fragments of Q' with the channel permutation described at the top: a0 / a2 hold channels 16j + 4t + {0, 1} / + // {2, 3} of row r0, a1 / a3 the same of row r0 + 8. + uint32_t qf[16][4]; +#pragma unroll + for (int j = 0; j < 16; ++j) { + const uint2 lo = *reinterpret_cast(sQ + r0 * kInt8QStride + 16 * j + 4 * t); + const uint2 hi = *reinterpret_cast(sQ + (r0 + 8) * kInt8QStride + 16 * j + 4 * t); + qf[j][0] = lo.x; + qf[j][1] = hi.x; + qf[j][2] = lo.y; + qf[j][3] = hi.y; + } + __syncthreads(); // sQ aliases sK8[0] + + auto resolve_index = [&](int raw) -> int { + return (raw >= 0 && raw < total && raw <= query_position && raw < capacity) ? raw : -1; + }; + + const int tiles = (count + kTile - 1) / kTile; + if ((tid >> 6) < tiles) { + const int tile = tid >> 6; + const int candidate = tile * kTile + (tid & 63); + sIdx[tile * kTile + (tid & 63)] = candidate < count ? resolve_index(row_indices[candidate]) : -1; + } + __syncthreads(); + + const uint32_t k_smem_base = static_cast(__cvta_generic_to_shared(sK8)); + const uint32_t v_smem_base = static_cast(__cvta_generic_to_shared(sV8)); + const uint32_t p_smem_base = static_cast(__cvta_generic_to_shared(sP)); + const int load_chunk = tid & 15; + + // 16 B per lane: 16 lanes per candidate row, 8 rows per pass, 8 passes per tile (K and V). + auto issue_tile = [&](int tile, int stage, int idx_buf) { +#pragma unroll + for (int i = 0; i < kTile / 8; ++i) { + const int cand = i * 8 + (tid >> 4); + const int index = sIdx[idx_buf * kTile + cand]; + const bool valid = index >= 0; + const int64_t source = static_cast(valid ? index : 0) * kHeadDim + load_chunk * 16; + const uint32_t offset = static_cast(stage * kInt8TileBytes + cand * kInt8RowStride + load_chunk * 16); + CpAsync16Zfill(k_smem_base + offset, key_base + source, valid); + CpAsync16Zfill(v_smem_base + offset, value_base + source, valid); + } + }; + + if (tiles > 0) { + issue_tile(0, 0, 0); + } + CpAsyncCommit(); + + float acc_o[8][4]; +#pragma unroll + for (int n = 0; n < 8; ++n) { + acc_o[n][0] = acc_o[n][1] = acc_o[n][2] = acc_o[n][3] = 0.0f; + } + // Softmax state of rows r0 / r0 + 8: the running maximum is identical in all threads of the row, the running sum only + // covers this thread's columns (merged across threads and warps after the loop). + float run_max0 = -FLT_MAX; + float run_max1 = -FLT_MAX; + float run_sum0 = 0.0f; + float run_sum1 = 0.0f; + const float scale_log2 = scale * 1.44269504088896341f; + + // Shared memory byte addresses that do not change from tile to tile. + const uint32_t k_lane_offset = + static_cast((warp * 16 + ((lane >> 3) & 1) * 8 + (lane & 7)) * kInt8RowStride + (lane >> 4) * 16); + const uint32_t v_lane_offset = + static_cast((((lane >> 3) & 1) * 8 + (lane & 7)) * kInt8RowStride + (warp * 4 + (lane >> 4)) * 16); + const uint32_t p_lane_offset = static_cast(((lane & 15) * kInt8PStride + (lane >> 4) * 8) * sizeof(__half)); + + int idx_cur = 0; // tile % 3 + int idx_next = 1 % kIdxStages; // (tile + 1) % 3 + int idx_next2 = 2 % kIdxStages; // (tile + 2) % 3 + for (int tile = 0; tile < tiles; ++tile) { + CpAsyncWait<0>(); // K / V of this tile have landed (this thread's copies) + __syncthreads(); // (1) everyone's copies are visible; the previous tile's smem consumers are done + + const int stage = tile & 1; + if (tile + 1 < tiles) { + issue_tile(tile + 1, stage ^ 1, idx_next); + } + CpAsyncCommit(); + // Raw indices of tile + 2: loaded now, resolved and stored at the end of the iteration (hides the load latency). + int raw_next2 = -1; + const bool stage_next2 = (tile + 2 < tiles) && tid < kTile; + if (stage_next2) { + const int candidate = (tile + 2) * kTile + tid; + raw_next2 = candidate < count ? row_indices[candidate] : -1; + } + + // S = Q' K^T for this warp's 16 candidates (two n tiles of 8). + float acc_s[2][4]; +#pragma unroll + for (int n = 0; n < 2; ++n) { + acc_s[n][0] = acc_s[n][1] = acc_s[n][2] = acc_s[n][3] = 0.0f; + } + { + const uint32_t k_addr = k_smem_base + stage * kInt8TileBytes + k_lane_offset; +#pragma unroll + for (int jj = 0; jj < 8; ++jj) { + uint32_t r[4]; + LdmatrixX4Addr(r, k_addr + jj * 32); + uint32_t b0; + uint32_t b1; + Int8x4ToHalf2Adjacent(r[0], b0, b1); + MmaF16(acc_s[0], qf[2 * jj], b0, b1); + Int8x4ToHalf2Adjacent(r[1], b0, b1); + MmaF16(acc_s[1], qf[2 * jj], b0, b1); + Int8x4ToHalf2Adjacent(r[2], b0, b1); + MmaF16(acc_s[0], qf[2 * jj + 1], b0, b1); + Int8x4ToHalf2Adjacent(r[3], b0, b1); + MmaF16(acc_s[1], qf[2 * jj + 1], b0, b1); + } + } + + // Scores in the log2 domain; invalid candidates are masked. + const int cbase = warp * 16 + 2 * t; + const int2 ia = *reinterpret_cast(sIdx + idx_cur * kTile + cbase); + const int2 ib = *reinterpret_cast(sIdx + idx_cur * kTile + cbase + 8); + const bool va0 = ia.x >= 0; + const bool va1 = ia.y >= 0; + const bool vb0 = ib.x >= 0; + const bool vb1 = ib.y >= 0; + const float s00 = va0 ? acc_s[0][0] * scale_log2 : -FLT_MAX; + const float s01 = va1 ? acc_s[0][1] * scale_log2 : -FLT_MAX; + const float s02 = vb0 ? acc_s[1][0] * scale_log2 : -FLT_MAX; + const float s03 = vb1 ? acc_s[1][1] * scale_log2 : -FLT_MAX; + const float s10 = va0 ? acc_s[0][2] * scale_log2 : -FLT_MAX; + const float s11 = va1 ? acc_s[0][3] * scale_log2 : -FLT_MAX; + const float s12 = vb0 ? acc_s[1][2] * scale_log2 : -FLT_MAX; + const float s13 = vb1 ? acc_s[1][3] * scale_log2 : -FLT_MAX; + float m0 = fmaxf(fmaxf(s00, s01), fmaxf(s02, s03)); + float m1 = fmaxf(fmaxf(s10, s11), fmaxf(s12, s13)); + m0 = fmaxf(m0, __shfl_xor_sync(0xffffffffu, m0, 1)); + m1 = fmaxf(m1, __shfl_xor_sync(0xffffffffu, m1, 1)); + m0 = fmaxf(m0, __shfl_xor_sync(0xffffffffu, m0, 2)); + m1 = fmaxf(m1, __shfl_xor_sync(0xffffffffu, m1, 2)); + if (t == 0) { + sM[r0 * 4 + warp] = m0; + sM[(r0 + 8) * 4 + warp] = m1; + } + __syncthreads(); // (2) the four warps' row maxima are visible + + const float4 mx0 = *reinterpret_cast(sM + r0 * 4); + const float4 mx1 = *reinterpret_cast(sM + (r0 + 8) * 4); + const float new_max0 = fmaxf(run_max0, fmaxf(fmaxf(mx0.x, mx0.y), fmaxf(mx0.z, mx0.w))); + const float new_max1 = fmaxf(run_max1, fmaxf(fmaxf(mx1.x, mx1.y), fmaxf(mx1.z, mx1.w))); + const float alpha0 = FastExp2(run_max0 - new_max0); + const float alpha1 = FastExp2(run_max1 - new_max1); + const float p00 = va0 ? FastExp2(s00 - new_max0) : 0.0f; + const float p01 = va1 ? FastExp2(s01 - new_max0) : 0.0f; + const float p02 = vb0 ? FastExp2(s02 - new_max0) : 0.0f; + const float p03 = vb1 ? FastExp2(s03 - new_max0) : 0.0f; + const float p10 = va0 ? FastExp2(s10 - new_max1) : 0.0f; + const float p11 = va1 ? FastExp2(s11 - new_max1) : 0.0f; + const float p12 = vb0 ? FastExp2(s12 - new_max1) : 0.0f; + const float p13 = vb1 ? FastExp2(s13 - new_max1) : 0.0f; + run_sum0 = run_sum0 * alpha0 + ((p00 + p01) + (p02 + p03)); + run_sum1 = run_sum1 * alpha1 + ((p10 + p11) + (p12 + p13)); + run_max0 = new_max0; + run_max1 = new_max1; + { + const __half2 h00 = __floats2half2_rn(p00, p01); + const __half2 h01 = __floats2half2_rn(p02, p03); + const __half2 h10 = __floats2half2_rn(p10, p11); + const __half2 h11 = __floats2half2_rn(p12, p13); + *reinterpret_cast<__half2*>(sP + r0 * kInt8PStride + cbase) = h00; + *reinterpret_cast<__half2*>(sP + r0 * kInt8PStride + cbase + 8) = h01; + *reinterpret_cast<__half2*>(sP + (r0 + 8) * kInt8PStride + cbase) = h10; + *reinterpret_cast<__half2*>(sP + (r0 + 8) * kInt8PStride + cbase + 8) = h11; + } + __syncthreads(); // (3) the probability tile is visible + + // The maximum rarely changes after the first tiles: skip the accumulator rescale when the whole warp is unchanged. + if (__any_sync(0xffffffffu, alpha0 != 1.0f || alpha1 != 1.0f)) { +#pragma unroll + for (int n = 0; n < 8; ++n) { + acc_o[n][0] *= alpha0; + acc_o[n][1] *= alpha0; + acc_o[n][2] *= alpha1; + acc_o[n][3] *= alpha1; + } + } + + // O += P V for this warp's 64 output channels (four 16-channel chunks, each an even and an odd n tile). + { + const uint32_t v_addr = v_smem_base + stage * kInt8TileBytes + v_lane_offset; +#pragma unroll + for (int kk = 0; kk < kTile / 16; ++kk) { + uint32_t a[4]; + LdmatrixX4Addr(a, p_smem_base + p_lane_offset + kk * 32); +#pragma unroll + for (int cp = 0; cp < 2; ++cp) { + uint32_t r[4]; + LdmatrixX4TransAddr(r, v_addr + kk * 16 * kInt8RowStride + cp * 32); + uint32_t e0; + uint32_t o0; + uint32_t e1; + uint32_t o1; + // r[0] / r[1]: candidates 0-7 / 8-15 of the k step, chunk 2cp; r[2] / r[3]: the same for chunk 2cp + 1. + Int8x4ToHalf2Strided(r[0], e0, o0); + Int8x4ToHalf2Strided(r[1], e1, o1); + MmaF16(acc_o[4 * cp], a, e0, e1); + MmaF16(acc_o[4 * cp + 1], a, o0, o1); + Int8x4ToHalf2Strided(r[2], e0, o0); + Int8x4ToHalf2Strided(r[3], e1, o1); + MmaF16(acc_o[4 * cp + 2], a, e0, e1); + MmaF16(acc_o[4 * cp + 3], a, o0, o1); + } + } + } + + if (stage_next2) { + sIdx[idx_next2 * kTile + tid] = resolve_index(raw_next2); + } + idx_cur = idx_next; + idx_next = idx_next2; + idx_next2 = idx_next2 + 1 == kIdxStages ? 0 : idx_next2 + 1; + } + + // Merge the per-thread running sums: across the four threads of a row, then across the four warps. + float l0 = run_sum0; + float l1 = run_sum1; + l0 += __shfl_xor_sync(0xffffffffu, l0, 1); + l1 += __shfl_xor_sync(0xffffffffu, l1, 1); + l0 += __shfl_xor_sync(0xffffffffu, l0, 2); + l1 += __shfl_xor_sync(0xffffffffu, l1, 2); + __syncthreads(); + if (t == 0) { + sM[r0 * 4 + warp] = l0; + sM[(r0 + 8) * 4 + warp] = l1; + } + __syncthreads(); + const float4 ls0 = *reinterpret_cast(sM + r0 * 4); + const float4 ls1 = *reinterpret_cast(sM + (r0 + 8) * 4); + const float sum0 = (ls0.x + ls0.y) + (ls0.z + ls0.w); + const float sum1 = (ls1.x + ls1.y) + (ls1.z + ls1.w); + const float inv0 = sum0 == 0.0f ? 0.0f : 1.0f / sum0; + const float inv1 = sum1 == 0.0f ? 0.0f : 1.0f / sum1; + + // Thread t owns channels 64 * warp + 16 * chunk + 4t + {0, 1, 2, 3} + // = (even tile c0, odd tile c0, even tile c1, odd tile c1). +#pragma unroll + for (int chunk = 0; chunk < 4; ++chunk) { + const int dim = warp * 64 + chunk * 16 + t * 4; + const float4 vs = *reinterpret_cast(value_scale + kv_head * kHeadDim + dim); + if (r0 < group) { + __nv_bfloat16* dst = output + (static_cast(row) * num_heads + kv_head * group + r0) * kHeadDim + dim; + const __nv_bfloat162 lo = + __floats2bfloat162_rn(acc_o[2 * chunk][0] * inv0 * vs.x, acc_o[2 * chunk + 1][0] * inv0 * vs.y); + const __nv_bfloat162 hi = + __floats2bfloat162_rn(acc_o[2 * chunk][1] * inv0 * vs.z, acc_o[2 * chunk + 1][1] * inv0 * vs.w); + uint2 packed; + packed.x = *reinterpret_cast(&lo); + packed.y = *reinterpret_cast(&hi); + *reinterpret_cast(dst) = packed; + } + if (r0 + 8 < group) { + __nv_bfloat16* dst = output + (static_cast(row) * num_heads + kv_head * group + r0 + 8) * kHeadDim + dim; + const __nv_bfloat162 lo = + __floats2bfloat162_rn(acc_o[2 * chunk][2] * inv1 * vs.x, acc_o[2 * chunk + 1][2] * inv1 * vs.y); + const __nv_bfloat162 hi = + __floats2bfloat162_rn(acc_o[2 * chunk][3] * inv1 * vs.z, acc_o[2 * chunk + 1][3] * inv1 * vs.w); + uint2 packed; + packed.x = *reinterpret_cast(&lo); + packed.y = *reinterpret_cast(&hi); + *reinterpret_cast(dst) = packed; + } + } +#else + (void)query; + (void)key; + (void)value; + (void)selected_indices; + (void)selected_counts; + (void)seqlens_k; + (void)output; + (void)sequence_length; + (void)num_heads; + (void)kv_num_heads; + (void)capacity; + (void)max_selected; + (void)scale; + (void)key_scale; + (void)value_scale; +#endif +} + +} // namespace dsa_prefill +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc index 003012a5a6a1b..c02629f55003c 100644 --- a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc @@ -593,14 +593,25 @@ TEST(DynamicSparseAttentionTest, GroupedDecodeRatio12ReferenceParity_CUDA) { } } -// Selected-only prefill without a softmax sink, the configuration the tensor-core BF16 prefill kernel serves. -DynamicSparseAttentionCase MakeSelectedOnlyPrefillCase(bool int8_cache) { - auto c = int8_cache ? MakeGroupedCacheCase(9, false, false, 2, 20, 2, true) - : MakeGroupedCacheCase(9, false, true, 2, 20, 2); +// Selected-only prefill without a softmax sink, the configuration the tensor-core prefill kernels serve. With +// invalid_entries every fifth selected index is replaced by one that must be skipped (negative, in the future of the +// query row, or outside the cache). +DynamicSparseAttentionCase MakeSelectedOnlyPrefillCase(bool int8_cache, int num_heads = 20, int selected_count = 509, + bool invalid_entries = false) { + auto c = int8_cache ? MakeGroupedCacheCase(9, false, false, 2, num_heads, 2, true, selected_count) + : MakeGroupedCacheCase(9, false, true, 2, num_heads, 2, false, selected_count); c.head_sink.clear(); c.smooth_softmax = 0; for (int row = 0; row < c.batch_size * c.sequence_length; ++row) { const int64_t b = row / c.sequence_length; + const int64_t position = c.seqlens_k[b] + 1 - c.sequence_length + row % c.sequence_length; + if (invalid_entries) { + const int32_t invalid_values[] = {-1, static_cast(position) + 1, + static_cast(c.cache_sequence_length) + 5, -7, 1 << 30}; + for (int i = 3; i < c.selected_counts[row]; i += 5) { + c.selected_indices[row * c.max_selected + i] = invalid_values[(i / 5) % 5]; + } + } for (int h = 0; h < c.num_heads; ++h) { const int64_t kv = h / (c.num_heads / c.kv_num_heads); const int64_t out = (row * c.num_heads + h) * c.head_size; @@ -608,6 +619,7 @@ DynamicSparseAttentionCase MakeSelectedOnlyPrefillCase(bool int8_cache) { double denominator = 0.0; for (int i = 0; i < c.selected_counts[row]; ++i) { const int index = c.selected_indices[row * c.max_selected + i]; + if (index < 0 || index > position || index >= c.cache_sequence_length) continue; const int64_t offset = ((b * c.kv_num_heads + kv) * c.cache_sequence_length + index) * c.head_size; double dot = 0.0; for (int d = 0; d < c.head_size; ++d) { @@ -641,6 +653,30 @@ TEST(DynamicSparseAttentionTest, Int8CachePrefillBFloat16_CUDA) { } } +// The tensor-core INT8 prefill kernel (ORT_DSA_PREFILL_INT8_TC=1) and the BF16-dequantizing one (=0) over query-head +// group sizes 1, 6, 12 and 16 (padded to the 16-row MMA tile), candidate counts around the 64-candidate tile and +// invalid selected indices. +TEST(DynamicSparseAttentionTest, Int8CachePrefillTensorCoreKernel_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + for (int num_heads : {2, 12, 24, 32}) { + for (int selected_count : {1, 64, 65, 509}) { + for (bool invalid_entries : {false, true}) { + SCOPED_TRACE("num_heads=" + std::to_string(num_heads) + ", selected_count=" + std::to_string(selected_count) + + ", invalid_entries=" + std::to_string(invalid_entries)); + const auto c = MakeSelectedOnlyPrefillCase(true, num_heads, selected_count, invalid_entries); + for (const char* int8_tc_env : {"1", "0"}) { + SCOPED_TRACE(std::string("ORT_DSA_PREFILL_INT8_TC=") + int8_tc_env); + ScopedEnvironmentVariables scoped_env_vars{ + EnvVarMap{{"ORT_DSA_PREFILL_KERNEL", "1"}, {"ORT_DSA_PREFILL_INT8_TC", int8_tc_env}}}; + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider()); + } + if (::testing::Test::HasFailure()) return; + } + } + } +} + TEST(DynamicSparseAttentionTest, FloatingCacheGroupedSingleSplitWorkspace_CUDA) { for (bool bf16 : {false, true}) { auto cuda_ep = DefaultCudaExecutionProvider(); @@ -2242,6 +2278,11 @@ std::vector MakeOffsetBufferCases(bool int8_cache, bool auxili cases.push_back(MisalignedBuffer("key_scale", &BufferOffsets::key_scale, scale_bytes)); cases.push_back(MisalignedBuffer("value_scale", &BufferOffsets::value_scale, scale_bytes)); all.key_scale = all.value_scale = scale_bytes; + // 8-byte cache and 4-byte output alignment still serve the BF16-dequantizing prefill kernel, but not the 16-byte + // cache copies and 8-byte output stores of the tensor-core INT8 prefill kernel. + cases.push_back(MisalignedBuffer("key_cache_8_byte", &BufferOffsets::key_cache, 8)); + cases.push_back(MisalignedBuffer("value_cache_8_byte", &BufferOffsets::value_cache, 8)); + cases.push_back(MisalignedBuffer("output_4_byte", &BufferOffsets::output, 4)); } if (auxiliary) { cases.push_back(MisalignedBuffer("auxiliary_key", &BufferOffsets::auxiliary_key, half_bytes)); @@ -2493,7 +2534,8 @@ TEST(DynamicSparseAttentionTest, Int8CacheGroupedDecodeMisalignedBuffers_CUDA) { TEST(DynamicSparseAttentionTest, Int8CachePrefillMisalignedBuffers_CUDA) { if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; - // Aligned buffers use the tensor-core prefill kernel; misaligned ones must fall back to the generic kernels. + // 16-byte cache / 8-byte output aligned buffers use the tensor-core INT8 prefill kernel, 8-byte / 4-byte aligned ones + // the BF16-dequantizing kernel and misaligned ones must fall back to the generic kernels. ExpectMisalignedBuffersMatch(MakeSelectedOnlyPrefillCase(true), true); } From 3c7c8c407e5d050046867b89bfc41441f8581b1c Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 29 Oct 2026 00:10:53 +0000 Subject: [PATCH 03/14] [CUDA] DynamicSparseAttention: range-safe INT8 prefill and session-option kill switches - INT8-cache prefill: when Q * key_scale leaves the FP16 range, a CTA-uniform FP32 path computes the tile instead of clamping, which corrupted valid scores (for example, cancelling large terms). The in-range fast path is unchanged. - Add the ep.cuda.dsa_prefill_int8_tc and ep.cuda.dsa_prefill_kernel session options (contrib_ops/cuda/cuda_session_options.h). They are resolved once at kernel construction, override ORT_DSA_PREFILL_INT8_TC and ORT_DSA_PREFILL_KERNEL, and reject invalid values, including an empty string, when the session is created. - Tests: position-dependent random K/V in the selected-only prefill sweeps (identical rows hid score and softmax errors), cancelling-pair and beyond-half-range query regressions, and a session-option test (PrefillSessionOptions_CUDA). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../onnxruntime_session_options_config_keys.h | 10 + .../cuda/bert/dynamic_sparse_attention.cc | 6 + .../cuda/bert/dynamic_sparse_attention.h | 3 + .../bert/dynamic_sparse_attention_impl.cu | 16 +- .../cuda/bert/dynamic_sparse_attention_impl.h | 2 + .../dynamic_sparse_attention_prefill_int8.cuh | 89 +++++++- .../contrib_ops/cuda/cuda_session_options.h | 38 ++++ .../dynamic_sparse_attention_op_test.cc | 196 ++++++++++++++++-- 8 files changed, 334 insertions(+), 26 deletions(-) create mode 100644 onnxruntime/contrib_ops/cuda/cuda_session_options.h diff --git a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h index 861a3f40ac1ba..610f31aa172c0 100644 --- a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h +++ b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h @@ -503,6 +503,16 @@ static const char* const kOrtSessionOptionsCudaFpAIntBGemvWaveAware = "ep.cuda.f static const char* const kOrtSessionOptionsCudaGqaWorkspaceMaxTotalSequenceLength = "ep.cuda.gqa_workspace_max_total_sequence_length"; +/// CUDA DynamicSparseAttention kernel switches. Each is resolved once when the kernel is created, so a change needs a new +/// session, and overrides the environment variable named below. The 0/1 switches accept only "0" and "1", and the +/// thresholds a non-negative decimal integer; any other value, including an empty string, fails session initialization. +/// "0" makes the INT8-cache BF16 prefill dequantize the gathered K/V tiles to BF16 instead of running FP16 tensor-core +/// MMA on them. Default "1". Overrides ORT_DSA_PREFILL_INT8_TC. +static const char* const kOrtSessionOptionsCudaDsaPrefillInt8Tc = "ep.cuda.dsa_prefill_int8_tc"; + +/// "0" routes the BF16 prefill through the generic kernels. Default "1". Overrides ORT_DSA_PREFILL_KERNEL. +static const char* const kOrtSessionOptionsCudaDsaPrefillKernel = "ep.cuda.dsa_prefill_kernel"; + /// Reader-only, hardware-neutral kernel dispatch intent shared across operators and EPs. /// Valid values are "auto", "latency", "memory", and "safe"; matching is case-sensitive. /// Unset or empty values map to "auto". Other values cause INVALID_ARGUMENT during diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc index d3e020eb16e2e..b1d2007e86a31 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc @@ -105,6 +105,10 @@ DynamicSparseAttention::DynamicSparseAttention(const OpKernelInfo& info) rotary_interleaved_ = ParseBoolAttribute(info, "rotary_interleaved", 0); use_smooth_softmax_ = ParseBoolAttribute(info, "smooth_softmax", 0); auxiliary_kv_shared_ = ParseBoolAttribute(info, "auxiliary_kv_shared", 0); + prefill_kernel_ = ReadCudaSessionOption(info, kOrtSessionOptionsCudaDsaPrefillKernel, + "ORT_DSA_PREFILL_KERNEL", 1, 0, 1) != 0; + prefill_int8_tc_ = ReadCudaSessionOption(info, kOrtSessionOptionsCudaDsaPrefillInt8Tc, + "ORT_DSA_PREFILL_INT8_TC", 1, 0, 1) != 0; strict_validation_ = ParseEnvironmentVariableWithDefault(kStrictValidation, false); attention_mode_ = ParseAttentionMode( info.GetAttrOrDefault("attention_mode", "selected_only")); @@ -203,6 +207,8 @@ Status DynamicSparseAttention::ComputeInternal(OpKernelContext* context) cons using CudaT = typename OrtToCudaType::type; DynamicSparseAttentionData data; + data.prefill_kernel = prefill_kernel_; + data.prefill_int8_tc = prefill_int8_tc_; data.query = reinterpret_cast(query->Data()); data.key = key == nullptr ? nullptr : reinterpret_cast(key->Data()); data.value = value == nullptr ? nullptr : reinterpret_cast(value->Data()); diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.h b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.h index 8719af8b3cb04..d14e0183f987c 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.h +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.h @@ -4,6 +4,7 @@ #pragma once #include "core/providers/cuda/cuda_kernel.h" +#include "contrib_ops/cuda/cuda_session_options.h" #include "contrib_ops/cpu/bert/dynamic_sparse_attention_helper.h" namespace onnxruntime { @@ -29,6 +30,8 @@ class DynamicSparseAttention final : public onnxruntime::cuda::CudaKernel { bool use_smooth_softmax_; bool auxiliary_kv_shared_; bool strict_validation_; + bool prefill_kernel_; + bool prefill_int8_tc_; DynamicSparseAttentionMode attention_mode_; DynamicSparseAttentionKvSource selected_kv_source_; }; diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu index 88d738f026a56..700f0fb0cb996 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu @@ -2077,7 +2077,7 @@ Status LaunchDynamicSparseAttention( bool use_int8_tc_prefill = false; if constexpr (std::is_same_v) { if (!use_grouped_decode) { - const bool enabled = ParseEnvironmentVariableWithDefault("ORT_DSA_PREFILL_KERNEL", 1) != 0; + const bool enabled = data.prefill_kernel; // max_shared_memory_per_block is the default 48 KB limit; the kernel opts in to more and requires SM80+. int device = 0; int cc_major = 0; @@ -2103,13 +2103,25 @@ Status LaunchDynamicSparseAttention( parameters.num_heads % parameters.kv_num_heads == 0 && parameters.num_heads / parameters.kv_num_heads <= dsa_prefill::kMaxGroup && static_cast(optin_shared) >= dsa_prefill::kSharedBytes; + if (use_prefill_kernel) { + cudaFuncAttributes attributes{}; + const auto kernel = data.quantized_cache ? dsa_prefill::DynamicSparseAttentionPrefillKernel + : dsa_prefill::DynamicSparseAttentionPrefillKernel; + CUDA_RETURN_IF_ERROR(cudaFuncGetAttributes(&attributes, kernel)); + use_prefill_kernel = attributes.ptxVersion >= 80; + } // The tensor-core INT8 variant converts the cached rows in registers; it needs 16-byte cache rows and 8-byte // output vectors (the BF16-dequantizing kernel above is the fallback). use_int8_tc_prefill = use_prefill_kernel && data.quantized_cache && - ParseEnvironmentVariableWithDefault("ORT_DSA_PREFILL_INT8_TC", 1) != 0 && + data.prefill_int8_tc && IsAligned(data.present_key, 16) && IsAligned(data.present_value, 16) && IsAligned(data.output, 8) && static_cast(optin_shared) >= dsa_prefill::kInt8SharedBytes; + if (use_int8_tc_prefill) { + cudaFuncAttributes attributes{}; + CUDA_RETURN_IF_ERROR(cudaFuncGetAttributes(&attributes, dsa_prefill::DynamicSparseAttentionPrefillInt8Kernel)); + use_int8_tc_prefill = attributes.ptxVersion >= 80; + } } } if (use_grouped_decode) { diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h index baa7e566eec17..eec2bd1a171fc 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h @@ -50,6 +50,8 @@ struct DynamicSparseAttentionData { T* output = nullptr; float* attention_workspace = nullptr; bool quantized_cache = false; + bool prefill_kernel = true; + bool prefill_int8_tc = true; const float* key_scale = nullptr; const float* value_scale = nullptr; }; diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh index dfbdd7ae8afc7..de0b5a9998476 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh @@ -7,6 +7,7 @@ // // The scales of the INT8 cache are per (kv head, channel) constants, so they leave the tile loop: // * the key scale is folded into the query once per CTA: Q' = Q * key_scale (rounded once to fp16), S = Q' K_int8; +// a CTA whose Q' does not fit fp16 (or is not finite) evaluates its query heads in fp32 instead (AttendRowFp32); // * the value scale is applied to the output once per CTA: O = (P V_int8) * value_scale. // The gathered int8 rows stay int8 in shared memory (cp.async, double buffered, one tile ahead). A fragment read with // ldmatrix is converted to fp16 exactly in registers (1024 + (byte ^ 0x80) - 1152, "magic number" conversion) and @@ -100,6 +101,65 @@ __device__ __forceinline__ void Int8x4ToHalf2Strided(uint32_t w, uint32_t& even, even = Half2MinusBias(a); odd = Half2MinusBias(b); } + +// FP32 evaluation of the query heads of one (row, KV head), one warp per head and eight channels per lane, for rows +// whose Q * key_scale leaves the FP16 range of the tensor-core path. Clipping there would break the cancellation +// between large terms. `query` and `output` point at the first head of the group; `key_scale` and `value_scale` at the +// KV head. +__device__ __noinline__ void AttendRowFp32(const __nv_bfloat16* __restrict__ query, const int8_t* __restrict__ key, + const int8_t* __restrict__ value, const float* __restrict__ key_scale, + const float* __restrict__ value_scale, + const int32_t* __restrict__ row_indices, int count, int total, + int query_position, float scale, int group, + __nv_bfloat16* __restrict__ output, int warp, int lane) { + float ks[8]; +#pragma unroll + for (int j = 0; j < 8; ++j) { + ks[j] = key_scale[lane + j * 32]; + } + for (int h = warp; h < group; h += kThreads / 32) { + float q[8]; +#pragma unroll + for (int j = 0; j < 8; ++j) { + q[j] = __bfloat162float(query[static_cast(h) * kHeadDim + lane + j * 32]); + } + float numerator[8] = {}; + float maximum = -FLT_MAX; + float denominator = 0.0f; + for (int i = 0; i < count; ++i) { + const int index = row_indices[i]; + if (index < 0 || index >= total || index > query_position) { + continue; + } + const int64_t offset = static_cast(index) * kHeadDim + lane; + float dot = 0.0f; +#pragma unroll + for (int j = 0; j < 8; ++j) { + dot += q[j] * (static_cast(key[offset + j * 32]) * ks[j]); + } +#pragma unroll + for (int delta = 16; delta > 0; delta /= 2) { + dot += __shfl_xor_sync(0xffffffffu, dot, delta); + } + const float score = dot * scale; + const float next_max = fmaxf(maximum, score); + const float correction = expf(maximum - next_max); + const float probability = expf(score - next_max); + denominator = denominator * correction + probability; +#pragma unroll + for (int j = 0; j < 8; ++j) { + numerator[j] = numerator[j] * correction + probability * static_cast(value[offset + j * 32]); + } + maximum = next_max; + } + const float inverse = denominator == 0.0f ? 0.0f : 1.0f / denominator; +#pragma unroll + for (int j = 0; j < 8; ++j) { + const int d = lane + j * 32; + output[static_cast(h) * kHeadDim + d] = __float2bfloat16_rn(numerator[j] * inverse * value_scale[d]); + } + } +} #endif // grid = (batch * sequence_length, kv_num_heads), block = kThreads, dynamic smem = kInt8SharedBytes. @@ -151,13 +211,19 @@ __global__ void __launch_bounds__(kThreads, 3) DynamicSparseAttentionPrefillInt8 const int8_t* key_base = key + kv_head_offset; const int8_t* value_base = value + kv_head_offset; - // Q' = Q * key_scale as fp16 (clamped so that it cannot overflow); the padding rows are zero so their (ignored) - // scores stay finite. + // Q' = Q * key_scale as fp16; the padding rows are zero so their (ignored) scores stay finite. A row whose Q' does + // not fit fp16 (including inf / NaN) is evaluated in fp32 by the whole CTA instead. Underflow to fp16 subnormals + // costs at most 2^-25 per element, below the rounding of the fp16 operands, so it stays on this path. + const __nv_bfloat16* query_group = query + (static_cast(row) * num_heads + kv_head * group) * kHeadDim; + bool needs_fp32 = false; { const int chunk = tid & 31; const float4 k0 = *reinterpret_cast(key_scale + kv_head * kHeadDim + chunk * 8); const float4 k1 = *reinterpret_cast(key_scale + kv_head * kHeadDim + chunk * 8 + 4); - const __nv_bfloat16* query_group = query + (static_cast(row) * num_heads + kv_head * group) * kHeadDim; + auto to_half2 = [&needs_fp32](float x, float y) { + needs_fp32 |= !(fabsf(x) <= 65504.0f) || !(fabsf(y) <= 65504.0f); + return __floats2half2_rn(x, y); + }; #pragma unroll for (int i = 0; i < kMaxGroup / 4; ++i) { const int r = (tid >> 5) + 4 * i; @@ -168,11 +234,10 @@ __global__ void __launch_bounds__(kThreads, 3) DynamicSparseAttentionPrefillInt8 const float2 q23 = __bfloat1622float2(*reinterpret_cast(&q.y)); const float2 q45 = __bfloat1622float2(*reinterpret_cast(&q.z)); const float2 q67 = __bfloat1622float2(*reinterpret_cast(&q.w)); - auto clamp = [](float x) { return fminf(fmaxf(x, -65504.0f), 65504.0f); }; - const __half2 h01 = __floats2half2_rn(clamp(q01.x * k0.x), clamp(q01.y * k0.y)); - const __half2 h23 = __floats2half2_rn(clamp(q23.x * k0.z), clamp(q23.y * k0.w)); - const __half2 h45 = __floats2half2_rn(clamp(q45.x * k1.x), clamp(q45.y * k1.y)); - const __half2 h67 = __floats2half2_rn(clamp(q67.x * k1.z), clamp(q67.y * k1.w)); + const __half2 h01 = to_half2(q01.x * k0.x, q01.y * k0.y); + const __half2 h23 = to_half2(q23.x * k0.z, q23.y * k0.w); + const __half2 h45 = to_half2(q45.x * k1.x, q45.y * k1.y); + const __half2 h67 = to_half2(q67.x * k1.z, q67.y * k1.w); packed.x = *reinterpret_cast(&h01); packed.y = *reinterpret_cast(&h23); packed.z = *reinterpret_cast(&h45); @@ -181,7 +246,13 @@ __global__ void __launch_bounds__(kThreads, 3) DynamicSparseAttentionPrefillInt8 *reinterpret_cast(sQ + r * kInt8QStride + chunk * 8) = packed; } } - __syncthreads(); + if (__syncthreads_or(needs_fp32)) { + // The condition is uniform across the CTA and no cp.async is in flight yet. + AttendRowFp32(query_group, key_base, value_base, key_scale + kv_head * kHeadDim, value_scale + kv_head * kHeadDim, + row_indices, count, total, query_position, scale, group, + output + (static_cast(row) * num_heads + kv_head * group) * kHeadDim, warp, lane); + return; + } // A fragments of Q' with the channel permutation described at the top: a0 / a2 hold channels 16j + 4t + {0, 1} / // {2, 3} of row r0, a1 / a3 the same of row r0 + 8. diff --git a/onnxruntime/contrib_ops/cuda/cuda_session_options.h b/onnxruntime/contrib_ops/cuda/cuda_session_options.h new file mode 100644 index 0000000000000..93f6ffcb64876 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/cuda_session_options.h @@ -0,0 +1,38 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include +#include +#include "core/providers/cuda/cuda_kernel.h" +#include "core/platform/env_var_utils.h" +#include "core/session/onnxruntime_session_options_config_keys.h" + +namespace onnxruntime::contrib::cuda { + +// An explicit configured value takes precedence over the environment. Throws unless the result is an integer in +// [min_value, max_value]. +inline int ResolveCudaSessionOption(const std::optional& configured, const char* key, const char* env, + int default_value, int min_value, int max_value) { + const std::string value = configured.has_value() + ? *configured + : ParseEnvironmentVariableWithDefault(env, std::to_string(default_value)); + int parsed = 0; + const auto result = std::from_chars(value.data(), value.data() + value.size(), parsed); + ORT_ENFORCE(result.ec == std::errc{} && result.ptr == value.data() + value.size() && + parsed >= min_value && parsed <= max_value, + key, " (", env, ") must be an integer in [", min_value, ", ", max_value, "], got '", value, "'."); + return parsed; +} + +// Call only from kernel construction. Explicit session configuration takes precedence over the environment. +inline int ReadCudaSessionOption(const OpKernelInfo& info, const char* key, const char* env, + int default_value, int min_value, int max_value) { + return ResolveCudaSessionOption(info.GetConfigOptions().GetConfigEntry(key), key, env, default_value, min_value, + max_value); +} + +} // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc index c02629f55003c..1cf8f40af7b35 100644 --- a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc @@ -8,6 +8,7 @@ #include #include #include +#include #include #include #include @@ -31,6 +32,7 @@ #ifdef USE_CUDA #include "core/framework/to_tensor_proto_element_type.h" #include "core/providers/cuda/cuda_provider_options.h" +#include "core/session/onnxruntime_session_options_config_keys.h" #include "test/common/cuda_op_test_utils.h" // Uses std::system, which is unavailable on iOS, so only CUDA builds (Windows/Linux) include it. #include "test/contrib_ops/fresh_process_test_util.h" @@ -120,7 +122,8 @@ void RunDynamicSparseAttentionCase( const DynamicSparseAttentionCase& c, std::unique_ptr execution_provider, OpTester::ExpectResult expected_result = OpTester::ExpectResult::kExpectSuccess, - const std::string& expected_error = "") { + const std::string& expected_error = "", + const std::vector>& session_configs = {}) { ASSERT_NE(execution_provider, nullptr); const int64_t hidden_size = c.num_heads * c.head_size; @@ -267,7 +270,15 @@ void RunDynamicSparseAttentionCase( std::vector> execution_providers; execution_providers.push_back(std::move(execution_provider)); - tester.Run(expected_result, expected_error, {}, nullptr, &execution_providers); + // The options of the default OpTester::Run overload plus the requested session configuration entries. + SessionOptions session_options; + session_options.use_per_session_threads = false; + session_options.session_log_verbosity_level = 1; + session_options.graph_optimization_level = TransformerLevel::Default; + for (const auto& [key, value] : session_configs) { + ASSERT_STATUS_OK(session_options.config_options.AddConfigEntry(key, value)); + } + tester.Run(session_options, expected_result, expected_error, {}, nullptr, &execution_providers); } DynamicSparseAttentionCase MakeSingleTokenSelectedOnlyCase() { @@ -596,22 +607,10 @@ TEST(DynamicSparseAttentionTest, GroupedDecodeRatio12ReferenceParity_CUDA) { // Selected-only prefill without a softmax sink, the configuration the tensor-core prefill kernels serve. With // invalid_entries every fifth selected index is replaced by one that must be skipped (negative, in the future of the // query row, or outside the cache). -DynamicSparseAttentionCase MakeSelectedOnlyPrefillCase(bool int8_cache, int num_heads = 20, int selected_count = 509, - bool invalid_entries = false) { - auto c = int8_cache ? MakeGroupedCacheCase(9, false, false, 2, num_heads, 2, true, selected_count) - : MakeGroupedCacheCase(9, false, true, 2, num_heads, 2, false, selected_count); - c.head_sink.clear(); - c.smooth_softmax = 0; +void ComputeSelectedOnlyPrefillExpectedOutput(DynamicSparseAttentionCase& c, bool int8_cache) { for (int row = 0; row < c.batch_size * c.sequence_length; ++row) { const int64_t b = row / c.sequence_length; const int64_t position = c.seqlens_k[b] + 1 - c.sequence_length + row % c.sequence_length; - if (invalid_entries) { - const int32_t invalid_values[] = {-1, static_cast(position) + 1, - static_cast(c.cache_sequence_length) + 5, -7, 1 << 30}; - for (int i = 3; i < c.selected_counts[row]; i += 5) { - c.selected_indices[row * c.max_selected + i] = invalid_values[(i / 5) % 5]; - } - } for (int h = 0; h < c.num_heads; ++h) { const int64_t kv = h / (c.num_heads / c.kv_num_heads); const int64_t out = (row * c.num_heads + h) * c.head_size; @@ -638,6 +637,48 @@ DynamicSparseAttentionCase MakeSelectedOnlyPrefillCase(bool int8_cache, int num_ } } } +} + +DynamicSparseAttentionCase MakeSelectedOnlyPrefillCase(bool int8_cache, int num_heads = 20, int selected_count = 509, + bool invalid_entries = false) { + auto c = int8_cache ? MakeGroupedCacheCase(9, false, false, 2, num_heads, 2, true, selected_count) + : MakeGroupedCacheCase(9, false, true, 2, num_heads, 2, false, selected_count); + c.head_sink.clear(); + c.smooth_softmax = 0; + if (int8_cache) { + // Distinct rows make wrong QK scores and online-softmax weights observable in the output; the larger scale + // spreads the softmax weights enough for score errors to show. The value scales shrink by 2^-4 (exact) so that + // the BF16 rounding of the dequantized operands in the BF16-dequantizing kernel, about half of the output + // tolerance at this magnitude, keeps a margin. + c.scale = 0.5f; + for (float& value_scale : c.value_scale) value_scale *= 0.0625f; + for (float& value : c.value) value *= 0.0625f; + std::mt19937 gen(17); + for (int b = 0; b < c.batch_size; ++b) { + const int past_length = c.seqlens_k[b] + 1 - c.sequence_length; + for (int kv = 0; kv < c.kv_num_heads; ++kv) { + for (int pos = 0; pos < past_length; ++pos) { + for (int d = 0; d < c.head_size; ++d) { + const int64_t offset = ((b * c.kv_num_heads + kv) * c.cache_sequence_length + pos) * c.head_size + d; + c.past_key[offset] = c.expected_present_key[offset] = static_cast(int(gen() % 256) - 128); + c.past_value[offset] = c.expected_present_value[offset] = static_cast(int(gen() % 256) - 128); + } + } + } + } + } + if (invalid_entries) { + for (int row = 0; row < c.batch_size * c.sequence_length; ++row) { + const int64_t b = row / c.sequence_length; + const int64_t position = c.seqlens_k[b] + 1 - c.sequence_length + row % c.sequence_length; + const int32_t invalid_values[] = {-1, static_cast(position) + 1, + static_cast(c.cache_sequence_length) + 5, -7, 1 << 30}; + for (int i = 3; i < c.selected_counts[row]; i += 5) { + c.selected_indices[row * c.max_selected + i] = invalid_values[(i / 5) % 5]; + } + } + } + ComputeSelectedOnlyPrefillExpectedOutput(c, int8_cache); return c; } @@ -677,6 +718,131 @@ TEST(DynamicSparseAttentionTest, Int8CachePrefillTensorCoreKernel_CUDA) { } } +std::vector> PrefillKernelConfigs(const char* kernel, const char* int8_tc) { + return {{kOrtSessionOptionsCudaDsaPrefillKernel, kernel}, {kOrtSessionOptionsCudaDsaPrefillInt8Tc, int8_tc}}; +} + +// A CTA of the tensor-core INT8 kernel whose Q * key_scale does not fit fp16 evaluates its scores in fp32. The query +// below cancels exactly: selected index 0 scores 131072 * 1 + 65536 * -2 = 0 like index 1 (all-zero key), so the output +// is the mean of the two values. Clamping Q * key_scale to the fp16 range would score index 0 at -65504 instead. +TEST(DynamicSparseAttentionTest, Int8PrefillCancellingLargeQuery_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + auto c = MakeSelectedOnlyPrefillCase(true, 2, 2); + std::fill(c.query.begin(), c.query.end(), 0.0f); + std::fill(c.key_scale.begin(), c.key_scale.end(), 1.0f); + std::fill(c.value_scale.begin(), c.value_scale.end(), 1.0f); + std::fill(c.key.begin(), c.key.end(), 0.0f); + std::fill(c.value.begin(), c.value.end(), 0.0f); + std::fill(c.past_key.begin(), c.past_key.end(), 0.0f); + std::fill(c.past_value.begin(), c.past_value.end(), 0.0f); + for (int b = 0; b < c.batch_size; ++b) { + for (int kv = 0; kv < c.kv_num_heads; ++kv) { + const int64_t offset = (b * c.kv_num_heads + kv) * c.cache_sequence_length * c.head_size; + c.past_key[offset] = 1.0f; + c.past_key[offset + 1] = -2.0f; + std::fill_n(c.past_value.begin() + offset, c.head_size, 2.0f); + std::fill_n(c.past_value.begin() + offset + c.head_size, c.head_size, 6.0f); + } + } + c.expected_present_key = c.past_key; + c.expected_present_value = c.past_value; + for (size_t i = 0; i < c.query.size(); i += c.head_size) { + c.query[i] = 131072.0f; + c.query[i + 1] = 65536.0f; + } + ComputeSelectedOnlyPrefillExpectedOutput(c, true); + for (size_t i = 0; i < c.expected_output.size(); i += c.head_size) { + const size_t row = i / (c.num_heads * c.head_size); + ASSERT_EQ(c.expected_output[i], row == 0 ? 0.0f : 4.0f); + } + + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), + OpTester::ExpectResult::kExpectSuccess, "", + PrefillKernelConfigs("1", "1")); +} + +// Multiplying every query by 2^28 and dividing the scale by 2^28 keeps the scores while pushing Q * key_scale of every +// CTA past fp16, so the fp32 path serves all query-head group sizes, candidate counts and invalid indices. A second +// case puts a huge query on a channel whose keys are all zero for a single head, so only some CTAs take the fp32 path. +TEST(DynamicSparseAttentionTest, Int8PrefillQueryBeyondHalfRange_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + const auto configs = PrefillKernelConfigs("1", "1"); + + for (int num_heads : {2, 12, 24, 32}) { + for (int selected_count : {1, 65, 509}) { + SCOPED_TRACE("num_heads=" + std::to_string(num_heads) + ", selected_count=" + std::to_string(selected_count)); + auto c = MakeSelectedOnlyPrefillCase(true, num_heads, selected_count, true); + for (auto& q : c.query) q *= 0x1p28f; + c.scale *= 0x1p-28f; + ComputeSelectedOnlyPrefillExpectedOutput(c, true); + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), + OpTester::ExpectResult::kExpectSuccess, "", configs); + if (::testing::Test::HasFailure()) return; + } + } + + auto c = MakeSelectedOnlyPrefillCase(true, 24, 65, true); + const int64_t channels = c.kv_num_heads * c.head_size; + constexpr int64_t kZeroChannel = 7; + // Every key of KV head 0 is zero on kZeroChannel, so a huge query there contributes exactly nothing to the scores. + for (int b = 0; b < c.batch_size; ++b) { + for (int64_t pos = 0; pos < c.cache_sequence_length; ++pos) { + const int64_t offset = ((b * c.kv_num_heads) * c.cache_sequence_length + pos) * c.head_size + kZeroChannel; + c.past_key[offset] = c.expected_present_key[offset] = 0.0f; + } + } + for (int row = 0; row < c.batch_size * c.sequence_length; ++row) { + c.key[row * channels + kZeroChannel] = 0.0f; + const int b = row / c.sequence_length; + const int pos = c.seqlens_k[b] + 1 - c.sequence_length + row % c.sequence_length; + c.expected_present_key[((b * c.kv_num_heads) * c.cache_sequence_length + pos) * c.head_size + kZeroChannel] = 0.0f; + if (row % 2 == 1) c.query[(row * c.num_heads + 5) * c.head_size + kZeroChannel] = 1.0e7f; + } + ComputeSelectedOnlyPrefillExpectedOutput(c, true); + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), + OpTester::ExpectResult::kExpectSuccess, "", configs); +} + +// Explicit ep.cuda.* entries take precedence over the environment, which is consulted only when no entry is set; an +// invalid value from either source fails kernel construction. +TEST(DynamicSparseAttentionTest, PrefillSessionOptions_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + const auto c = MakeSelectedOnlyPrefillCase(true, 12, 65, true); + const auto expect_success = OpTester::ExpectResult::kExpectSuccess; + const auto expect_failure = OpTester::ExpectResult::kExpectFailure; + + { + ScopedEnvironmentVariables env{ + EnvVarMap{{"ORT_DSA_PREFILL_KERNEL", "invalid"}, {"ORT_DSA_PREFILL_INT8_TC", "invalid"}}}; + for (const char* kernel : {"0", "1"}) { + for (const char* int8_tc : {"0", "1"}) { + SCOPED_TRACE(std::string("kernel=") + kernel + ", int8_tc=" + int8_tc); + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_success, "", + PrefillKernelConfigs(kernel, int8_tc)); + } + } + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_failure, + "ORT_DSA_PREFILL_INT8_TC", + {{kOrtSessionOptionsCudaDsaPrefillKernel, "1"}}); + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_failure, + "ORT_DSA_PREFILL_KERNEL", + {{kOrtSessionOptionsCudaDsaPrefillInt8Tc, "1"}}); + } + for (const char* int8_tc : {"0", "1"}) { + ScopedEnvironmentVariables env{EnvVarMap{{"ORT_DSA_PREFILL_KERNEL", "1"}, {"ORT_DSA_PREFILL_INT8_TC", int8_tc}}}; + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_success); + } + for (const char* invalid : {"2", "-1", "", "1x", "true"}) { + SCOPED_TRACE(std::string("value='") + invalid + "'"); + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_failure, + kOrtSessionOptionsCudaDsaPrefillInt8Tc, + {{kOrtSessionOptionsCudaDsaPrefillInt8Tc, invalid}}); + } +} + TEST(DynamicSparseAttentionTest, FloatingCacheGroupedSingleSplitWorkspace_CUDA) { for (bool bf16 : {false, true}) { auto cuda_ep = DefaultCudaExecutionProvider(); From df87fb45cc7bbfaa000f279591693305760a0637 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Wed, 28 Oct 2026 22:19:56 +0000 Subject: [PATCH 04/14] [CUDA] Fix MSVC C4244 narrowing in the DSA INT8 prefill tests The test case stores its dimensions as int64_t, so the int locals initialized from them raised C4244, which the Windows CUDA build treats as an error. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../test/contrib_ops/dynamic_sparse_attention_op_test.cc | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc index 1cf8f40af7b35..a554732613058 100644 --- a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc @@ -655,7 +655,7 @@ DynamicSparseAttentionCase MakeSelectedOnlyPrefillCase(bool int8_cache, int num_ for (float& value : c.value) value *= 0.0625f; std::mt19937 gen(17); for (int b = 0; b < c.batch_size; ++b) { - const int past_length = c.seqlens_k[b] + 1 - c.sequence_length; + const int64_t past_length = c.seqlens_k[b] + 1 - c.sequence_length; for (int kv = 0; kv < c.kv_num_heads; ++kv) { for (int pos = 0; pos < past_length; ++pos) { for (int d = 0; d < c.head_size; ++d) { @@ -793,10 +793,10 @@ TEST(DynamicSparseAttentionTest, Int8PrefillQueryBeyondHalfRange_CUDA) { c.past_key[offset] = c.expected_present_key[offset] = 0.0f; } } - for (int row = 0; row < c.batch_size * c.sequence_length; ++row) { + for (int64_t row = 0; row < c.batch_size * c.sequence_length; ++row) { c.key[row * channels + kZeroChannel] = 0.0f; - const int b = row / c.sequence_length; - const int pos = c.seqlens_k[b] + 1 - c.sequence_length + row % c.sequence_length; + const int64_t b = row / c.sequence_length; + const int64_t pos = c.seqlens_k[b] + 1 - c.sequence_length + row % c.sequence_length; c.expected_present_key[((b * c.kv_num_heads) * c.cache_sequence_length + pos) * c.head_size + kZeroChannel] = 0.0f; if (row % 2 == 1) c.query[(row * c.num_heads + 5) * c.head_size + kZeroChannel] = 1.0e7f; } From 54d90aefcd575242b19d858b56f794c250ec9101 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 29 Oct 2026 00:11:16 +0000 Subject: [PATCH 05/14] [CUDA] DynamicSparseAttention: split-K tensor-core INT8 kernel for decode (SM 8.0+) The BF16 INT8-cache decode path (sequence_length <= 8, 12 query heads per KV head) ran a CUDA-core kernel that is instruction-issue / shared-load bound on the per-candidate serial chain (70 us per layer at 2K..32K context on H200, 6.9% of a speculative-decoding cycle of Qwen3.8-Flash-Next). DynamicSparseAttentionInt8TcKernel (the mma.sync m16n8k16 prefill kernel) is now templated on kSplitK. With kSplitK each CTA of the grid (row, KV head, split) consumes a contiguous range of 64-candidate tiles and writes the un-normalized partial output, the natural-log row max and the row sum in the layout of the grouped decode merge kernel, which is reused unchanged. The split count targets about 384 CTAs (one wave of three CTAs per SM), at most 32. Routing requires a BF16 INT8 cache, head_size 256, selected-only main KV source, no smooth softmax / sinks, 16-byte aligned buffers, the opt-in shared memory and PTX of at least compute_80; everything else keeps the old kernel. A CTA whose Q * key_scale does not fit fp16 (or is not finite) evaluates its range of candidates in fp32 and publishes an fp32 partial in the same layout, so the range-safe fallback of the prefill kernel also holds for the splits. ep.cuda.dsa_decode_int8_tc=0 (ORT_DSA_DECODE_INT8_TC) restores the old kernel bit for bit; ep.cuda.dsa_decode_int8_tc_tiles_per_split (ORT_DSA_DECODE_INT8_TC_TILES_PER_SPLIT) overrides the split heuristic (tuning and tests). Both follow the strict parsing of the other ep.cuda.dsa_* options and are resolved once when the kernel is created. H200, 8 rows, 2048 selected: 70.3 us + 5.3 us merge -> 9.5 us + 3.5 us merge per layer; decode -4.6 .. -6.8% ms/cycle at 2K / 8K / 32K on the three Qwen3.8-Flash-Next packages. The fp32 accumulation order differs from the CUDA-core kernel, so generated-token hashes are not identical; ep.cuda.dsa_decode_int8_tc=0 reproduces the old results exactly. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../onnxruntime_session_options_config_keys.h | 9 + .../cuda/bert/dynamic_sparse_attention.cc | 7 + .../cuda/bert/dynamic_sparse_attention.h | 2 + .../bert/dynamic_sparse_attention_impl.cu | 96 ++++++++- .../cuda/bert/dynamic_sparse_attention_impl.h | 2 + .../dynamic_sparse_attention_prefill_int8.cuh | 187 ++++++++++++++---- .../dynamic_sparse_attention_op_test.cc | 154 ++++++++++++++- 7 files changed, 401 insertions(+), 56 deletions(-) diff --git a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h index 610f31aa172c0..367b4db8ab1e3 100644 --- a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h +++ b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h @@ -513,6 +513,15 @@ static const char* const kOrtSessionOptionsCudaDsaPrefillInt8Tc = "ep.cuda.dsa_p /// "0" routes the BF16 prefill through the generic kernels. Default "1". Overrides ORT_DSA_PREFILL_KERNEL. static const char* const kOrtSessionOptionsCudaDsaPrefillKernel = "ep.cuda.dsa_prefill_kernel"; +/// "0" makes the INT8-cache BF16 decode (up to 8 query rows) use the CUDA-core grouped kernel instead of the split-K +/// tensor-core kernel. Default "1". Overrides ORT_DSA_DECODE_INT8_TC. +static const char* const kOrtSessionOptionsCudaDsaDecodeInt8Tc = "ep.cuda.dsa_decode_int8_tc"; + +/// Number of 64-candidate tiles each CTA of the split-K tensor-core decode kernel consumes. "0" derives it from the +/// problem size (about 384 CTAs in total). Default "0". Overrides ORT_DSA_DECODE_INT8_TC_TILES_PER_SPLIT. +static const char* const kOrtSessionOptionsCudaDsaDecodeInt8TcTilesPerSplit = + "ep.cuda.dsa_decode_int8_tc_tiles_per_split"; + /// Reader-only, hardware-neutral kernel dispatch intent shared across operators and EPs. /// Valid values are "auto", "latency", "memory", and "safe"; matching is case-sensitive. /// Unset or empty values map to "auto". Other values cause INVALID_ARGUMENT during diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc index b1d2007e86a31..9bd7e62b99008 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc @@ -109,6 +109,11 @@ DynamicSparseAttention::DynamicSparseAttention(const OpKernelInfo& info) "ORT_DSA_PREFILL_KERNEL", 1, 0, 1) != 0; prefill_int8_tc_ = ReadCudaSessionOption(info, kOrtSessionOptionsCudaDsaPrefillInt8Tc, "ORT_DSA_PREFILL_INT8_TC", 1, 0, 1) != 0; + decode_int8_tc_ = ReadCudaSessionOption(info, kOrtSessionOptionsCudaDsaDecodeInt8Tc, + "ORT_DSA_DECODE_INT8_TC", 1, 0, 1) != 0; + decode_int8_tc_tiles_per_split_ = + ReadCudaSessionOption(info, kOrtSessionOptionsCudaDsaDecodeInt8TcTilesPerSplit, + "ORT_DSA_DECODE_INT8_TC_TILES_PER_SPLIT", 0, 0, std::numeric_limits::max()); strict_validation_ = ParseEnvironmentVariableWithDefault(kStrictValidation, false); attention_mode_ = ParseAttentionMode( info.GetAttrOrDefault("attention_mode", "selected_only")); @@ -209,6 +214,8 @@ Status DynamicSparseAttention::ComputeInternal(OpKernelContext* context) cons DynamicSparseAttentionData data; data.prefill_kernel = prefill_kernel_; data.prefill_int8_tc = prefill_int8_tc_; + data.decode_int8_tc = decode_int8_tc_; + data.decode_int8_tc_tiles_per_split = decode_int8_tc_tiles_per_split_; data.query = reinterpret_cast(query->Data()); data.key = key == nullptr ? nullptr : reinterpret_cast(key->Data()); data.value = value == nullptr ? nullptr : reinterpret_cast(value->Data()); diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.h b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.h index d14e0183f987c..40f8d78d58859 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.h +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.h @@ -32,6 +32,8 @@ class DynamicSparseAttention final : public onnxruntime::cuda::CudaKernel { bool strict_validation_; bool prefill_kernel_; bool prefill_int8_tc_; + bool decode_int8_tc_; + int decode_int8_tc_tiles_per_split_; DynamicSparseAttentionMode attention_mode_; DynamicSparseAttentionKvSource selected_kv_source_; }; diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu index 700f0fb0cb996..ce82f26c38877 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu @@ -40,6 +40,7 @@ constexpr int kGroupedDecodeCandidatesPerSplit = 64; constexpr int kTargetGroupedDecodeBlocks = 128; constexpr int kGroupedQuantizedDecodeCandidatesPerSplit = 32; constexpr int kTargetGroupedQuantizedDecodeBlocks = 512; +constexpr int kTargetInt8TcDecodeCtas = 384; template __device__ __forceinline__ float ToFloat(T value) { @@ -1807,6 +1808,24 @@ int GetGroupedDecodeSplitCount(const DynamicSparseAttentionParameters& parameter return static_cast(std::min({candidate_splits, occupancy_splits, max_splits})); } +// Split count of the tensor-core INT8 decode kernel: one CTA per (row, KV head, split) and about +// kTargetInt8TcDecodeCtas CTAs in total (one wave of three CTAs per SM on 132 SMs), each consuming at least one +// 64-candidate tile. The time per CTA is dominated by a fixed latency chain, so the kernel is flat in the tiles per +// split (H200, 8 rows, 33 tiles: 1 -> 10.8 us, 2 -> 9.5, 3 -> 10.7, 4 -> 9.9, 6 -> 11.9). It never exceeds the split +// count the workspace was sized for. A positive tiles_per_split_override replaces the automatic choice. +int GetInt8TcDecodeSplitCount(const DynamicSparseAttentionParameters& parameters, int workspace_splits, + int tiles_per_split_override) { + const int64_t tiles = (static_cast(parameters.max_selected) + dsa_prefill::kTile - 1) / dsa_prefill::kTile; + const int64_t row_kv_heads = + static_cast(parameters.batch_size) * parameters.sequence_length * parameters.kv_num_heads; + const int64_t tiles_per_split = + tiles_per_split_override > 0 + ? tiles_per_split_override + : std::max(1, (tiles * row_kv_heads + kTargetInt8TcDecodeCtas - 1) / kTargetInt8TcDecodeCtas); + const int64_t splits = (tiles + tiles_per_split - 1) / tiles_per_split; + return static_cast(std::max(1, std::min({splits, workspace_splits, 32}))); +} + size_t GetValidationHashEntries(const DynamicSparseAttentionParameters& parameters) { if (parameters.max_selected == 0) { return 0; @@ -2119,13 +2138,48 @@ Status LaunchDynamicSparseAttention( static_cast(optin_shared) >= dsa_prefill::kInt8SharedBytes; if (use_int8_tc_prefill) { cudaFuncAttributes attributes{}; - CUDA_RETURN_IF_ERROR(cudaFuncGetAttributes(&attributes, dsa_prefill::DynamicSparseAttentionPrefillInt8Kernel)); + CUDA_RETURN_IF_ERROR(cudaFuncGetAttributes(&attributes, dsa_prefill::DynamicSparseAttentionInt8TcKernel)); use_int8_tc_prefill = attributes.ptxVersion >= 80; } } } + // Decode (sequence_length <= 8) over the INT8 cache: the same tensor-core pipeline split over the candidate tiles of + // each (row, KV head); the partials go through the grouped decode merge kernel. + bool use_int8_tc_decode = false; + if constexpr (std::is_same_v) { + if (use_grouped_decode && data.quantized_cache && data.decode_int8_tc) { + int device = 0; + int cc_major = 0; + int optin_shared = 0; + CUDA_RETURN_IF_ERROR(cudaGetDevice(&device)); + CUDA_RETURN_IF_ERROR(cudaDeviceGetAttribute(&cc_major, cudaDevAttrComputeCapabilityMajor, device)); + CUDA_RETURN_IF_ERROR(cudaDeviceGetAttribute(&optin_shared, cudaDevAttrMaxSharedMemoryPerBlockOptin, device)); + constexpr int kMaxGridDimY = 65535; + // The kernel loads 16-byte query / scale vectors, gathers 16-byte cache chunks and stores 16-byte partials; + // buffers bound through IOBinding can start at any element offset, so misaligned ones use the generic kernel. + use_int8_tc_decode = + cc_major >= 8 && parameters.head_size == dsa_prefill::kHeadDim && !local_plus_selected && + parameters.selected_kv_source == DynamicSparseAttentionKvSource::kMain && + !parameters.use_smooth_softmax && parameters.max_selected > 0 && parameters.kv_num_heads <= kMaxGridDimY && + parameters.num_heads / parameters.kv_num_heads <= dsa_prefill::kMaxGroup && + IsAligned(data.prepared_query, 16) && IsAligned(data.present_key, 16) && IsAligned(data.present_value, 16) && + IsAligned(data.key_scale, 16) && IsAligned(data.value_scale, 16) && + IsAligned(data.attention_workspace, 16) && + static_cast(optin_shared) >= dsa_prefill::kInt8SharedBytes; + if (use_int8_tc_decode) { + // A build whose PTX is older than compute_80 holds an empty stub for the kernel. + cudaFuncAttributes attributes{}; + CUDA_RETURN_IF_ERROR(cudaFuncGetAttributes(&attributes, dsa_prefill::DynamicSparseAttentionInt8TcKernel)); + use_int8_tc_decode = attributes.ptxVersion >= 80; + } + } + } if (use_grouped_decode) { - const int split_count = GetGroupedDecodeSplitCount(parameters, data.quantized_cache); + const int workspace_splits = GetGroupedDecodeSplitCount(parameters, data.quantized_cache); + const int split_count = + use_int8_tc_decode + ? GetInt8TcDecodeSplitCount(parameters, workspace_splits, data.decode_int8_tc_tiles_per_split) + : workspace_splits; const int64_t partial_blocks = query_blocks * split_count; const int64_t grouped_blocks = static_cast(parameters.batch_size) * parameters.sequence_length * parameters.kv_num_heads * @@ -2137,9 +2191,30 @@ Status LaunchDynamicSparseAttention( float* partial_max = data.attention_workspace; float* partial_sum = partial_max + partial_blocks; float* partial_output = partial_sum + partial_blocks; + if (use_int8_tc_decode) { + // The partial vectors come first so that they stay 16-byte aligned whatever the number of partial blocks. + partial_output = data.attention_workspace; + partial_max = partial_output + partial_blocks * parameters.head_size; + partial_sum = partial_max + partial_blocks; + } const size_t grouped_shared_bytes = GetGroupedAttentionSharedBytes(parameters, data.quantized_cache ? sizeof(float) : sizeof(T)); - if (data.quantized_cache) { + if (use_int8_tc_decode) { + if constexpr (std::is_same_v) { + CUDA_RETURN_IF_ERROR(cudaFuncSetAttribute( + dsa_prefill::DynamicSparseAttentionInt8TcKernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(dsa_prefill::kInt8SharedBytes))); + const dim3 grid(static_cast(parameters.batch_size * parameters.sequence_length), + static_cast(parameters.kv_num_heads), static_cast(split_count)); + dsa_prefill::DynamicSparseAttentionInt8TcKernel + <<>>( + data.prepared_query, reinterpret_cast(data.present_key), + reinterpret_cast(data.present_value), data.selected_indices, data.selected_counts, + data.seqlens_k, nullptr, partial_max, partial_sum, partial_output, split_count, + parameters.sequence_length, parameters.num_heads, parameters.kv_num_heads, parameters.cache_capacity, + parameters.max_selected, parameters.scale, data.key_scale, data.value_scale); + } + } else if (data.quantized_cache) { static const int vectorized_staging = ParseEnvironmentVariableWithDefault("ORT_DSA_INT8_VECTOR_STAGING", 1); ORT_ENFORCE(vectorized_staging == 0 || vectorized_staging == 1, @@ -2272,14 +2347,15 @@ Status LaunchDynamicSparseAttention( static_cast(parameters.kv_num_heads)); if (use_int8_tc_prefill) { CUDA_RETURN_IF_ERROR(cudaFuncSetAttribute( - dsa_prefill::DynamicSparseAttentionPrefillInt8Kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + dsa_prefill::DynamicSparseAttentionInt8TcKernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(dsa_prefill::kInt8SharedBytes))); - dsa_prefill::DynamicSparseAttentionPrefillInt8Kernel<<>>( - data.prepared_query, reinterpret_cast(data.present_key), - reinterpret_cast(data.present_value), data.selected_indices, data.selected_counts, - data.seqlens_k, data.output, parameters.sequence_length, parameters.num_heads, parameters.kv_num_heads, - parameters.cache_capacity, parameters.max_selected, parameters.scale, data.key_scale, data.value_scale); + dsa_prefill::DynamicSparseAttentionInt8TcKernel + <<>>( + data.prepared_query, reinterpret_cast(data.present_key), + reinterpret_cast(data.present_value), data.selected_indices, data.selected_counts, + data.seqlens_k, data.output, nullptr, nullptr, nullptr, 1, parameters.sequence_length, + parameters.num_heads, parameters.kv_num_heads, parameters.cache_capacity, parameters.max_selected, + parameters.scale, data.key_scale, data.value_scale); } else if (data.quantized_cache) { CUDA_RETURN_IF_ERROR(cudaFuncSetAttribute( dsa_prefill::DynamicSparseAttentionPrefillKernel, cudaFuncAttributeMaxDynamicSharedMemorySize, diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h index eec2bd1a171fc..24f96686efbf4 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h @@ -52,6 +52,8 @@ struct DynamicSparseAttentionData { bool quantized_cache = false; bool prefill_kernel = true; bool prefill_int8_tc = true; + bool decode_int8_tc = true; + int decode_int8_tc_tiles_per_split = 0; const float* key_scale = nullptr; const float* value_scale = nullptr; }; diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh index de0b5a9998476..4a518dc618619 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh @@ -5,6 +5,12 @@ // bfloat16 query / output, head_size 256). One CTA handles one (query row, KV head); the heads of the GQA group form // the M dimension of an m16n8k16 MMA (padded to 16). // +// The same body also serves decode (kSplitK = true): the grid gets a third dimension, each CTA consumes a contiguous +// range of the row's 64-candidate tiles and writes an un-normalized partial result (running maximum in the natural +// logarithm domain, running sum, value-scaled accumulator) in the layout of the grouped decode kernel, which the +// existing merge kernel combines. This replaces the CUDA-core grouped decode kernel, whose per-candidate dot product, +// warp reduction and exponentials made it issue-bound, with the tensor-core pipeline below. +// // The scales of the INT8 cache are per (kv head, channel) constants, so they leave the tile loop: // * the key scale is folded into the query once per CTA: Q' = Q * key_scale (rounded once to fp16), S = Q' K_int8; // a CTA whose Q' does not fit fp16 (or is not finite) evaluates its query heads in fp32 instead (AttendRowFp32); @@ -105,13 +111,19 @@ __device__ __forceinline__ void Int8x4ToHalf2Strided(uint32_t w, uint32_t& even, // FP32 evaluation of the query heads of one (row, KV head), one warp per head and eight channels per lane, for rows // whose Q * key_scale leaves the FP16 range of the tensor-core path. Clipping there would break the cancellation // between large terms. `query` and `output` point at the first head of the group; `key_scale` and `value_scale` at the -// KV head. +// KV head. The candidates [begin, end) of `row_indices` are consumed: all of them for prefill, the range of the CTA +// for a decode split. With kSplitK the un-normalized partial (running maximum in the natural logarithm domain, running +// sum, value-scaled accumulator) of block partial_block + head * split_count goes to partial_*, as the main path does; +// otherwise the normalized result goes to output. +template __device__ __noinline__ void AttendRowFp32(const __nv_bfloat16* __restrict__ query, const int8_t* __restrict__ key, const int8_t* __restrict__ value, const float* __restrict__ key_scale, const float* __restrict__ value_scale, - const int32_t* __restrict__ row_indices, int count, int total, + const int32_t* __restrict__ row_indices, int begin, int end, int total, int query_position, float scale, int group, - __nv_bfloat16* __restrict__ output, int warp, int lane) { + __nv_bfloat16* __restrict__ output, float* __restrict__ partial_max, + float* __restrict__ partial_sum, float* __restrict__ partial_output, + int partial_block, int split_count, int warp, int lane) { float ks[8]; #pragma unroll for (int j = 0; j < 8; ++j) { @@ -126,7 +138,7 @@ __device__ __noinline__ void AttendRowFp32(const __nv_bfloat16* __restrict__ que float numerator[8] = {}; float maximum = -FLT_MAX; float denominator = 0.0f; - for (int i = 0; i < count; ++i) { + for (int i = begin; i < end; ++i) { const int index = row_indices[i]; if (index < 0 || index >= total || index > query_position) { continue; @@ -152,20 +164,37 @@ __device__ __noinline__ void AttendRowFp32(const __nv_bfloat16* __restrict__ que } maximum = next_max; } - const float inverse = denominator == 0.0f ? 0.0f : 1.0f / denominator; + if constexpr (kSplitK) { + const int block = partial_block + h * split_count; + if (lane == 0) { + partial_max[block] = maximum; + partial_sum[block] = denominator; + } #pragma unroll - for (int j = 0; j < 8; ++j) { - const int d = lane + j * 32; - output[static_cast(h) * kHeadDim + d] = __float2bfloat16_rn(numerator[j] * inverse * value_scale[d]); + for (int j = 0; j < 8; ++j) { + const int d = lane + j * 32; + partial_output[static_cast(block) * kHeadDim + d] = numerator[j] * value_scale[d]; + } + } else { + const float inverse = denominator == 0.0f ? 0.0f : 1.0f / denominator; +#pragma unroll + for (int j = 0; j < 8; ++j) { + const int d = lane + j * 32; + output[static_cast(h) * kHeadDim + d] = __float2bfloat16_rn(numerator[j] * inverse * value_scale[d]); + } } } } #endif -// grid = (batch * sequence_length, kv_num_heads), block = kThreads, dynamic smem = kInt8SharedBytes. +// grid = (batch * sequence_length, kv_num_heads) or, with kSplitK, (batch * sequence_length, kv_num_heads, +// split_count); block = kThreads, dynamic smem = kInt8SharedBytes. // query: [rows, num_heads, 256] (already normalized / rotated); key, value: int8 [batch, kv_heads, capacity, 256]; -// key_scale, value_scale: float [kv_heads, 256]. -__global__ void __launch_bounds__(kThreads, 3) DynamicSparseAttentionPrefillInt8Kernel( +// key_scale, value_scale: float [kv_heads, 256]. The normalized result goes to output; with kSplitK the partials +// go to partial_max / partial_sum [rows * num_heads * split_count] and partial_output [.., 256] instead, indexed +// (row * num_heads + head) * split_count + split (output is unused; partial_output must be 16-byte aligned). +template +__global__ void __launch_bounds__(kThreads, 3) DynamicSparseAttentionInt8TcKernel( const __nv_bfloat16* __restrict__ query, const int8_t* __restrict__ key, const int8_t* __restrict__ value, @@ -173,6 +202,10 @@ __global__ void __launch_bounds__(kThreads, 3) DynamicSparseAttentionPrefillInt8 const int32_t* __restrict__ selected_counts, const int32_t* __restrict__ seqlens_k, __nv_bfloat16* __restrict__ output, + float* __restrict__ partial_max, + float* __restrict__ partial_sum, + float* __restrict__ partial_output, + int split_count, int sequence_length, int num_heads, int kv_num_heads, @@ -211,6 +244,26 @@ __global__ void __launch_bounds__(kThreads, 3) DynamicSparseAttentionPrefillInt8 const int8_t* key_base = key + kv_head_offset; const int8_t* value_base = value + kv_head_offset; + // The 64-candidate tiles this CTA consumes: all of the row's for prefill, a contiguous range for a decode split. + // A split without tiles only publishes an empty partial (maximum -FLT_MAX, sum 0), which the merge skips. + int tile_begin = 0; + int tiles = (count + kTile - 1) / kTile; + if constexpr (kSplitK) { + const int split = static_cast(blockIdx.z); + const int tiles_per_split = (tiles + split_count - 1) / split_count; + const int first = min(tiles, split * tiles_per_split); + tiles = min(tiles, first + tiles_per_split) - first; + tile_begin = first; + if (tiles == 0) { + if (tid < group) { + const int partial_block = (row * num_heads + kv_head * group + tid) * split_count + split; + partial_max[partial_block] = -FLT_MAX; + partial_sum[partial_block] = 0.0f; + } + return; + } + } + // Q' = Q * key_scale as fp16; the padding rows are zero so their (ignored) scores stay finite. A row whose Q' does // not fit fp16 (including inf / NaN) is evaluated in fp32 by the whole CTA instead. Underflow to fp16 subnormals // costs at most 2^-25 per element, below the rounding of the fp16 operands, so it stays on this path. @@ -248,9 +301,20 @@ __global__ void __launch_bounds__(kThreads, 3) DynamicSparseAttentionPrefillInt8 } if (__syncthreads_or(needs_fp32)) { // The condition is uniform across the CTA and no cp.async is in flight yet. - AttendRowFp32(query_group, key_base, value_base, key_scale + kv_head * kHeadDim, value_scale + kv_head * kHeadDim, - row_indices, count, total, query_position, scale, group, - output + (static_cast(row) * num_heads + kv_head * group) * kHeadDim, warp, lane); + const float* kv_key_scale = key_scale + kv_head * kHeadDim; + const float* kv_value_scale = value_scale + kv_head * kHeadDim; + if constexpr (kSplitK) { + const int partial_block = (row * num_heads + kv_head * group) * split_count + static_cast(blockIdx.z); + AttendRowFp32(query_group, key_base, value_base, kv_key_scale, kv_value_scale, row_indices, + tile_begin * kTile, min(count, (tile_begin + tiles) * kTile), total, query_position, scale, + group, nullptr, partial_max, partial_sum, partial_output, partial_block, split_count, warp, + lane); + } else { + AttendRowFp32(query_group, key_base, value_base, kv_key_scale, kv_value_scale, row_indices, 0, count, + total, query_position, scale, group, + output + (static_cast(row) * num_heads + kv_head * group) * kHeadDim, nullptr, + nullptr, nullptr, 0, 1, warp, lane); + } return; } @@ -272,10 +336,10 @@ __global__ void __launch_bounds__(kThreads, 3) DynamicSparseAttentionPrefillInt8 return (raw >= 0 && raw < total && raw <= query_position && raw < capacity) ? raw : -1; }; - const int tiles = (count + kTile - 1) / kTile; + const int first_candidate = tile_begin * kTile; if ((tid >> 6) < tiles) { const int tile = tid >> 6; - const int candidate = tile * kTile + (tid & 63); + const int candidate = first_candidate + tile * kTile + (tid & 63); sIdx[tile * kTile + (tid & 63)] = candidate < count ? resolve_index(row_indices[candidate]) : -1; } __syncthreads(); @@ -340,7 +404,7 @@ __global__ void __launch_bounds__(kThreads, 3) DynamicSparseAttentionPrefillInt8 int raw_next2 = -1; const bool stage_next2 = (tile + 2 < tiles) && tid < kTile; if (stage_next2) { - const int candidate = (tile + 2) * kTile + tid; + const int candidate = first_candidate + (tile + 2) * kTile + tid; raw_next2 = candidate < count ? row_indices[candidate] : -1; } @@ -494,33 +558,68 @@ __global__ void __launch_bounds__(kThreads, 3) DynamicSparseAttentionPrefillInt8 const float inv0 = sum0 == 0.0f ? 0.0f : 1.0f / sum0; const float inv1 = sum1 == 0.0f ? 0.0f : 1.0f / sum1; - // Thread t owns channels 64 * warp + 16 * chunk + 4t + {0, 1, 2, 3} - // = (even tile c0, odd tile c0, even tile c1, odd tile c1). + if constexpr (kSplitK) { + // Un-normalized partial of this split in the natural logarithm domain, as the merge kernel expects. Thread t owns + // channels 64 * warp + 16 * chunk + 4t + {0, 1, 2, 3} (see below), i.e. one float4 per chunk and row. + constexpr float kLn2 = 0.693147180559945309f; + const int split = static_cast(blockIdx.z); + const int partial_block0 = (row * num_heads + kv_head * group + r0) * split_count + split; + const int partial_block1 = partial_block0 + 8 * split_count; + if (warp == 0 && t == 0) { + if (r0 < group) { + partial_max[partial_block0] = run_max0 * kLn2; + partial_sum[partial_block0] = sum0; + } + if (r0 + 8 < group) { + partial_max[partial_block1] = run_max1 * kLn2; + partial_sum[partial_block1] = sum1; + } + } #pragma unroll - for (int chunk = 0; chunk < 4; ++chunk) { - const int dim = warp * 64 + chunk * 16 + t * 4; - const float4 vs = *reinterpret_cast(value_scale + kv_head * kHeadDim + dim); - if (r0 < group) { - __nv_bfloat16* dst = output + (static_cast(row) * num_heads + kv_head * group + r0) * kHeadDim + dim; - const __nv_bfloat162 lo = - __floats2bfloat162_rn(acc_o[2 * chunk][0] * inv0 * vs.x, acc_o[2 * chunk + 1][0] * inv0 * vs.y); - const __nv_bfloat162 hi = - __floats2bfloat162_rn(acc_o[2 * chunk][1] * inv0 * vs.z, acc_o[2 * chunk + 1][1] * inv0 * vs.w); - uint2 packed; - packed.x = *reinterpret_cast(&lo); - packed.y = *reinterpret_cast(&hi); - *reinterpret_cast(dst) = packed; + for (int chunk = 0; chunk < 4; ++chunk) { + const int dim = warp * 64 + chunk * 16 + t * 4; + const float4 vs = *reinterpret_cast(value_scale + kv_head * kHeadDim + dim); + if (r0 < group) { + *reinterpret_cast(partial_output + static_cast(partial_block0) * kHeadDim + dim) = + make_float4(acc_o[2 * chunk][0] * vs.x, acc_o[2 * chunk + 1][0] * vs.y, acc_o[2 * chunk][1] * vs.z, + acc_o[2 * chunk + 1][1] * vs.w); + } + if (r0 + 8 < group) { + *reinterpret_cast(partial_output + static_cast(partial_block1) * kHeadDim + dim) = + make_float4(acc_o[2 * chunk][2] * vs.x, acc_o[2 * chunk + 1][2] * vs.y, acc_o[2 * chunk][3] * vs.z, + acc_o[2 * chunk + 1][3] * vs.w); + } } - if (r0 + 8 < group) { - __nv_bfloat16* dst = output + (static_cast(row) * num_heads + kv_head * group + r0 + 8) * kHeadDim + dim; - const __nv_bfloat162 lo = - __floats2bfloat162_rn(acc_o[2 * chunk][2] * inv1 * vs.x, acc_o[2 * chunk + 1][2] * inv1 * vs.y); - const __nv_bfloat162 hi = - __floats2bfloat162_rn(acc_o[2 * chunk][3] * inv1 * vs.z, acc_o[2 * chunk + 1][3] * inv1 * vs.w); - uint2 packed; - packed.x = *reinterpret_cast(&lo); - packed.y = *reinterpret_cast(&hi); - *reinterpret_cast(dst) = packed; + } else { + // Thread t owns channels 64 * warp + 16 * chunk + 4t + {0, 1, 2, 3} + // = (even tile c0, odd tile c0, even tile c1, odd tile c1). +#pragma unroll + for (int chunk = 0; chunk < 4; ++chunk) { + const int dim = warp * 64 + chunk * 16 + t * 4; + const float4 vs = *reinterpret_cast(value_scale + kv_head * kHeadDim + dim); + if (r0 < group) { + __nv_bfloat16* dst = output + (static_cast(row) * num_heads + kv_head * group + r0) * kHeadDim + dim; + const __nv_bfloat162 lo = + __floats2bfloat162_rn(acc_o[2 * chunk][0] * inv0 * vs.x, acc_o[2 * chunk + 1][0] * inv0 * vs.y); + const __nv_bfloat162 hi = + __floats2bfloat162_rn(acc_o[2 * chunk][1] * inv0 * vs.z, acc_o[2 * chunk + 1][1] * inv0 * vs.w); + uint2 packed; + packed.x = *reinterpret_cast(&lo); + packed.y = *reinterpret_cast(&hi); + *reinterpret_cast(dst) = packed; + } + if (r0 + 8 < group) { + __nv_bfloat16* dst = + output + (static_cast(row) * num_heads + kv_head * group + r0 + 8) * kHeadDim + dim; + const __nv_bfloat162 lo = + __floats2bfloat162_rn(acc_o[2 * chunk][2] * inv1 * vs.x, acc_o[2 * chunk + 1][2] * inv1 * vs.y); + const __nv_bfloat162 hi = + __floats2bfloat162_rn(acc_o[2 * chunk][3] * inv1 * vs.z, acc_o[2 * chunk + 1][3] * inv1 * vs.w); + uint2 packed; + packed.x = *reinterpret_cast(&lo); + packed.y = *reinterpret_cast(&hi); + *reinterpret_cast(dst) = packed; + } } } #else @@ -531,6 +630,10 @@ __global__ void __launch_bounds__(kThreads, 3) DynamicSparseAttentionPrefillInt8 (void)selected_counts; (void)seqlens_k; (void)output; + (void)partial_max; + (void)partial_sum; + (void)partial_output; + (void)split_count; (void)sequence_length; (void)num_heads; (void)kv_num_heads; diff --git a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc index a554732613058..104d0582391c4 100644 --- a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc @@ -604,7 +604,7 @@ TEST(DynamicSparseAttentionTest, GroupedDecodeRatio12ReferenceParity_CUDA) { } } -// Selected-only prefill without a softmax sink, the configuration the tensor-core prefill kernels serve. With +// Selected-only attention without a softmax sink, the configuration the tensor-core kernels serve. With // invalid_entries every fifth selected index is replaced by one that must be skipped (negative, in the future of the // query row, or outside the cache). void ComputeSelectedOnlyPrefillExpectedOutput(DynamicSparseAttentionCase& c, bool int8_cache) { @@ -640,9 +640,9 @@ void ComputeSelectedOnlyPrefillExpectedOutput(DynamicSparseAttentionCase& c, boo } DynamicSparseAttentionCase MakeSelectedOnlyPrefillCase(bool int8_cache, int num_heads = 20, int selected_count = 509, - bool invalid_entries = false) { - auto c = int8_cache ? MakeGroupedCacheCase(9, false, false, 2, num_heads, 2, true, selected_count) - : MakeGroupedCacheCase(9, false, true, 2, num_heads, 2, false, selected_count); + bool invalid_entries = false, int sequence_length = 9) { + auto c = int8_cache ? MakeGroupedCacheCase(sequence_length, false, false, 2, num_heads, 2, true, selected_count) + : MakeGroupedCacheCase(sequence_length, false, true, 2, num_heads, 2, false, selected_count); c.head_sink.clear(); c.smooth_softmax = 0; if (int8_cache) { @@ -722,6 +722,12 @@ std::vector> PrefillKernelConfigs(const char return {{kOrtSessionOptionsCudaDsaPrefillKernel, kernel}, {kOrtSessionOptionsCudaDsaPrefillInt8Tc, int8_tc}}; } +std::vector> DecodeKernelConfigs(const char* int8_tc, + const char* tiles_per_split = "0") { + return {{kOrtSessionOptionsCudaDsaDecodeInt8Tc, int8_tc}, + {kOrtSessionOptionsCudaDsaDecodeInt8TcTilesPerSplit, tiles_per_split}}; +} + // A CTA of the tensor-core INT8 kernel whose Q * key_scale does not fit fp16 evaluates its scores in fp32. The query // below cancels exactly: selected index 0 scores 131072 * 1 + 65536 * -2 = 0 like index 1 (all-zero key), so the output // is the mean of the two values. Clamping Q * key_scale to the fp16 range would score index 0 at -65504 instead. @@ -843,6 +849,138 @@ TEST(DynamicSparseAttentionTest, PrefillSessionOptions_CUDA) { } } +// The split-K tensor-core INT8 decode kernel (ep.cuda.dsa_decode_int8_tc=1) and the CUDA-core grouped decode kernel +// (=0) over 1-8 query rows per batch item (row 0 of the case has no selected candidate), query-head group sizes 4, 6, +// 12 and 16 (padded to the 16-row MMA tile), candidate counts around the 64-candidate tile, invalid selected indices +// and 1, 3 or the automatic number of 64-candidate tiles per split (a split without tiles publishes an empty partial). +TEST(DynamicSparseAttentionTest, Int8CacheDecodeTensorCoreKernel_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + struct Config { + const char* decode_tc; + const char* tiles_per_split; + }; + const Config configs[] = {{"1", "0"}, {"1", "1"}, {"1", "3"}, {"0", "0"}}; + for (int sequence_length : {1, 3, 8}) { + for (int num_heads : {8, 12, 24, 32}) { + for (int selected_count : {1, 65, 509}) { + for (bool invalid_entries : {false, true}) { + SCOPED_TRACE("sequence_length=" + std::to_string(sequence_length) + + ", num_heads=" + std::to_string(num_heads) + + ", selected_count=" + std::to_string(selected_count) + + ", invalid_entries=" + std::to_string(invalid_entries)); + const auto c = MakeSelectedOnlyPrefillCase(true, num_heads, selected_count, invalid_entries, sequence_length); + for (const auto& config : configs) { + SCOPED_TRACE(std::string("dsa_decode_int8_tc=") + config.decode_tc + + ", dsa_decode_int8_tc_tiles_per_split=" + config.tiles_per_split); + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), + OpTester::ExpectResult::kExpectSuccess, "", + DecodeKernelConfigs(config.decode_tc, + config.tiles_per_split)); + } + if (::testing::Test::HasFailure()) return; + } + } + } + } +} + +// A decode CTA whose Q * key_scale does not fit fp16 evaluates its range of candidates in fp32 and publishes an fp32 +// partial for the merge kernel. Multiplying every query by 2^28 and dividing the scale by 2^28 keeps the scores while +// pushing every CTA past fp16, for any number of splits; a second case puts a huge query on a channel whose keys are +// all zero for a single head of every other row, so that only some CTAs take the fp32 path. +TEST(DynamicSparseAttentionTest, Int8DecodeQueryBeyondHalfRange_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + for (const char* tiles_per_split : {"0", "1", "3"}) { + const auto configs = DecodeKernelConfigs("1", tiles_per_split); + for (int sequence_length : {1, 3}) { + for (int num_heads : {8, 12, 32}) { + for (int selected_count : {1, 65, 509}) { + SCOPED_TRACE("tiles_per_split=" + std::string(tiles_per_split) + + ", sequence_length=" + std::to_string(sequence_length) + + ", num_heads=" + std::to_string(num_heads) + + ", selected_count=" + std::to_string(selected_count)); + auto c = MakeSelectedOnlyPrefillCase(true, num_heads, selected_count, true, sequence_length); + for (auto& q : c.query) q *= 0x1p28f; + c.scale *= 0x1p-28f; + ComputeSelectedOnlyPrefillExpectedOutput(c, true); + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), + OpTester::ExpectResult::kExpectSuccess, "", configs); + if (::testing::Test::HasFailure()) return; + } + } + } + + auto c = MakeSelectedOnlyPrefillCase(true, 24, 509, true, 3); + const int64_t channels = c.kv_num_heads * c.head_size; + constexpr int64_t kZeroChannel = 7; + for (int b = 0; b < c.batch_size; ++b) { + for (int64_t pos = 0; pos < c.cache_sequence_length; ++pos) { + const int64_t offset = ((b * c.kv_num_heads) * c.cache_sequence_length + pos) * c.head_size + kZeroChannel; + c.past_key[offset] = c.expected_present_key[offset] = 0.0f; + } + } + for (int64_t row = 0; row < c.batch_size * c.sequence_length; ++row) { + c.key[row * channels + kZeroChannel] = 0.0f; + const int64_t b = row / c.sequence_length; + const int64_t pos = c.seqlens_k[b] + 1 - c.sequence_length + row % c.sequence_length; + c.expected_present_key[((b * c.kv_num_heads) * c.cache_sequence_length + pos) * c.head_size + kZeroChannel] = + 0.0f; + if (row % 2 == 1) c.query[(row * c.num_heads + 5) * c.head_size + kZeroChannel] = 1.0e7f; + } + ComputeSelectedOnlyPrefillExpectedOutput(c, true); + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), + OpTester::ExpectResult::kExpectSuccess, "", configs); + if (::testing::Test::HasFailure()) return; + } +} + +// Explicit ep.cuda.* entries take precedence over the environment, which is consulted only when no entry is set; an +// invalid value from either source fails kernel construction. +TEST(DynamicSparseAttentionTest, DecodeSessionOptions_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + const auto c = MakeSelectedOnlyPrefillCase(true, 12, 65, true, 3); + const auto expect_success = OpTester::ExpectResult::kExpectSuccess; + const auto expect_failure = OpTester::ExpectResult::kExpectFailure; + + { + ScopedEnvironmentVariables env{ + EnvVarMap{{"ORT_DSA_DECODE_INT8_TC", "invalid"}, {"ORT_DSA_DECODE_INT8_TC_TILES_PER_SPLIT", "invalid"}}}; + for (const char* int8_tc : {"0", "1"}) { + for (const char* tiles_per_split : {"0", "2"}) { + SCOPED_TRACE(std::string("int8_tc=") + int8_tc + ", tiles_per_split=" + tiles_per_split); + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_success, "", + DecodeKernelConfigs(int8_tc, tiles_per_split)); + } + } + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_failure, + "(ORT_DSA_DECODE_INT8_TC_TILES_PER_SPLIT)", + {{kOrtSessionOptionsCudaDsaDecodeInt8Tc, "1"}}); + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_failure, + "(ORT_DSA_DECODE_INT8_TC)", + {{kOrtSessionOptionsCudaDsaDecodeInt8TcTilesPerSplit, "0"}}); + } + for (const char* int8_tc : {"0", "1"}) { + ScopedEnvironmentVariables env{ + EnvVarMap{{"ORT_DSA_DECODE_INT8_TC", int8_tc}, {"ORT_DSA_DECODE_INT8_TC_TILES_PER_SPLIT", "2"}}}; + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_success); + } + for (const char* invalid : {"2", "-1", "", "1x", "true"}) { + SCOPED_TRACE(std::string("value='") + invalid + "'"); + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_failure, + kOrtSessionOptionsCudaDsaDecodeInt8Tc, + {{kOrtSessionOptionsCudaDsaDecodeInt8Tc, invalid}}); + } + for (const char* invalid : {"-1", "", "1x", "many"}) { + SCOPED_TRACE(std::string("value='") + invalid + "'"); + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_failure, + kOrtSessionOptionsCudaDsaDecodeInt8TcTilesPerSplit, + {{kOrtSessionOptionsCudaDsaDecodeInt8TcTilesPerSplit, invalid}}); + } +} + TEST(DynamicSparseAttentionTest, FloatingCacheGroupedSingleSplitWorkspace_CUDA) { for (bool bf16 : {false, true}) { auto cuda_ep = DefaultCudaExecutionProvider(); @@ -2705,6 +2843,14 @@ TEST(DynamicSparseAttentionTest, Int8CachePrefillMisalignedBuffers_CUDA) { ExpectMisalignedBuffersMatch(MakeSelectedOnlyPrefillCase(true), true); } +TEST(DynamicSparseAttentionTest, Int8CacheDecodeMisalignedBuffers_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + // 16-byte aligned buffers use the split-K tensor-core INT8 decode kernel; misaligned ones must fall back to the + // CUDA-core grouped decode kernel. + ExpectMisalignedBuffersMatch(MakeSelectedOnlyPrefillCase(true, 24, 509, false, 3), true); +} + TEST(DynamicSparseAttentionTest, FloatingCacheGroupedDecodeMisalignedBuffers_CUDA) { if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; for (bool bf16 : {true, false}) { From e1de24c95f611acbc43b6bc94403e079f45d5cd7 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 29 Oct 2026 00:11:25 +0000 Subject: [PATCH 06/14] [CUDA] HyperConnectionLowRankMix: branch-free sigmoid (bit-exact, -26% PreMixVec8) SigmoidF32 chose between 1 / (1 + expf(-v)) and expf(v) / (1 + expf(v)) with a per-element branch. In the prefill pre-mix the signs of neighbouring elements differ, so nearly every warp executed both paths. Both paths call expf with -|v|, so use one expf and a select: e = expf(-|v|); result = (v >= 0 ? 1 : e) / (1 + e). The result is bit-identical for all 2^32 fp32 bit patterns (exhaustive check in a standalone harness, 0 mismatches; same nvcc flags, no fast-math). Measured (H200, hidden 2560, 4 branches, 8192-token chunk): PreMixVec8Kernel 112.1 -> 82.7 us per launch (96 launches per chunk, nsys), ScaledSiluKernel 7.3 -> 6.1 us; TTFT -0.3..-1.1% (all_int4 / top16 / top08 at 2K / 8K / 32K, 9 / 9 cells faster, first tokens identical). Decode stays within +-0.3%: the fused HcUpPreKernel calls the same function for at most 8 rows. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cuda/bert/hyper_connection_low_rank_mix_impl.cu | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.cu b/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.cu index 59c9626b0035b..c44181198969e 100644 --- a/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.cu @@ -51,10 +51,12 @@ constexpr int kMaxFusedHidden = kMaxVecIters * kThreads * 8; constexpr int kMaxFusedRank = 512; constexpr int kMaxFusedBranches = 8; +// expf(-|v|) is the argument both branches of the textbook form pass to expf (-v for v >= 0, v for v < 0), so one expf and +// one division give the bit-identical result (verified over all 2^32 fp32 bit patterns) without the divergent two-way +// branch, which executed both paths in nearly every warp of the prefill pre-mix. __device__ __forceinline__ float SigmoidF32(float v) { - if (v >= 0.0f) return 1.0f / (1.0f + expf(-v)); - const float e = expf(v); - return e / (1.0f + e); + const float e = expf(-fabsf(v)); + return (v >= 0.0f ? 1.0f : e) / (1.0f + e); } template From de38dc6789a14b1a37353c8db9f8bad52acf8e5a Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 29 Oct 2026 00:12:26 +0000 Subject: [PATCH 07/14] [CUDA] Document the INT8-KV DynamicSparseAttention tensor-core kernels, the early PDL trigger and the sigmoid change - dynamic_sparse_attention.md: the ep.cuda.dsa_* session options next to their environment variables (read once at kernel construction, strict 0/1 parsing, the tuning knob), the dispatch-table rows and alignment fallbacks of the INT8-cache prefill and split-K decode kernels, the FP32 range-safe fallback and its cost, and the H200 / Qwen3.8-Flash-Next numbers. - hyper_connection_low_rank_mix.md: the early griddepcontrol.launch_dependents trigger of ProjectedSigmoidGateMulKernel and the branch-free SigmoidF32 (bit-exact for all 2^32 float values). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cuda/dynamic_sparse_attention.md | 38 ++++++++++++++++++- .../cuda/hyper_connection_low_rank_mix.md | 21 ++++++++++ 2 files changed, 58 insertions(+), 1 deletion(-) diff --git a/docs/contrib_ops/cuda/dynamic_sparse_attention.md b/docs/contrib_ops/cuda/dynamic_sparse_attention.md index 86ca09455b0a3..dad9e1ed4d7d6 100644 --- a/docs/contrib_ops/cuda/dynamic_sparse_attention.md +++ b/docs/contrib_ops/cuda/dynamic_sparse_attention.md @@ -47,7 +47,14 @@ the code is only reported through `validation_status`. `ORT_DSA_INT8_VECTOR_STAGING` and `ORT_DSA_SHARED_SCALE_MERGE` are read once per process (cached in function-local statics) when the grouped decode path is first launched, so set them before the first operator invocation. Only `0` and `1` are accepted. Neither changes schemas or cache layouts. -`ORT_DSA_PREFILL_KERNEL` (default `1`) is read on every BF16 call. +`ORT_DSA_PREFILL_KERNEL`, `ORT_DSA_PREFILL_INT8_TC` and `ORT_DSA_DECODE_INT8_TC` (all default `1`) are read once +at kernel construction. Session options `ep.cuda.dsa_prefill_kernel`, `ep.cuda.dsa_prefill_int8_tc` and +`ep.cuda.dsa_decode_int8_tc` take precedence over their corresponding environment variables. +They accept only `0` or `1`; changing a value requires a new session. The loaded prefill and decode +specializations must have a PTX target of SM80+; otherwise dispatch uses the generic fallback. +`ORT_DSA_DECODE_INT8_TC_TILES_PER_SPLIT` (session option `ep.cuda.dsa_decode_int8_tc_tiles_per_split`) is a +tuning and test knob with the same precedence: an integer in `[0, INT_MAX]`, `0` keeps the automatic split +heuristic. Invalid values of any of these options fail session initialization. Grouped decode is used when `1 <= sequence_length <= 8`, `num_heads / kv_num_heads >= 4` (and divisible), `head_size == 256`, 16-bit `T` (FP16/BF16), more than 256 candidates @@ -63,6 +70,12 @@ and 8 otherwise. | | | | `0` | `MergeGroupedDecodeDynamicSparseAttentionKernel`: per-channel `expf` over all splits. | | `ORT_DSA_PREFILL_KERNEL` | `1` | BF16, not grouped decode, `sequence_length >= 8`, `head_size == 256`, `selected_only` from the main cache, `smooth_softmax = 0`, SM80+ with enough opt-in shared memory | `1` | `DynamicSparseAttentionPrefillKernel` tensor-core prefill. | | | | | `0` | Generic fused/split or grouped prefill kernels. | +| `ORT_DSA_PREFILL_INT8_TC` | `1` | The `ORT_DSA_PREFILL_KERNEL` conditions with an INT8 main cache, 16-byte-aligned `present_key` and `present_value`, an 8-byte-aligned `output` and enough opt-in shared memory for `dsa_prefill::kInt8SharedBytes` | `1` | `DynamicSparseAttentionInt8TcKernel` ([dynamic_sparse_attention_prefill_int8.cuh](../../../onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh)): the gathered K/V rows stay INT8 in shared memory and are converted exactly to FP16 in registers for FP16 `mma` with FP32 accumulation. `key_scale` is folded into the query once per block and `value_scale` is applied to the output once. A block whose scaled query does not fit FP16 (`|Q * key_scale| > 65504`, inf or NaN) is evaluated by an FP32 path (`AttendRowFp32`) instead of being clipped, which would break the cancellation between large terms. That path is about 48x slower for the affected blocks (8192 rows on an H200: 2.75 ms to 131 ms); set `ep.cuda.dsa_prefill_int8_tc=0` for models whose scaled queries exceed the FP16 range, because the kernel below has no such limit. | +| | | | `0` | `DynamicSparseAttentionPrefillKernel`: dequantizes every gathered tile to BF16 in shared memory. | +| `ORT_DSA_DECODE_INT8_TC` | `1` | Grouped decode (conditions above) with BF16, an INT8 main cache, `head_size == 256`, `selected_only` from the main cache (no `local_plus_selected`), `smooth_softmax = 0`, at most `kMaxGroup` query heads per KV head, 16-byte-aligned `prepared_query`, `present_key`, `present_value`, `key_scale`, `value_scale` and attention workspace, SM80+ and enough opt-in shared memory for `dsa_prefill::kInt8SharedBytes` | `1` | `DynamicSparseAttentionInt8TcKernel`: the tensor-core pipeline of the row above, split over the 64-candidate tiles of each (row, KV head). Each CTA of the (row, KV head, split) grid writes an un-normalized partial output, the natural-log row maximum and the row sum in the layout of the grouped decode merge kernel, which is reused unchanged. The split count targets 384 CTAs (one wave of three per SM, at most 32 splits). A CTA whose scaled query does not fit FP16 evaluates its candidates in FP32 (`AttendRowFp32`) and publishes an FP32 partial in the same layout, so the range-safe fallback also holds for the splits. | +| | | | `0` | The CUDA-core `GroupedDecodeDynamicSparseAttentionKernel`, bit for bit as before. | +| `ORT_DSA_DECODE_INT8_TC_TILES_PER_SPLIT` | `0` | `ORT_DSA_DECODE_INT8_TC=1` and the conditions above | `0` | Automatic: `max(1, ceil(tiles * rows * kv_heads / 384))` tiles per split, with `rows = batch_size * sequence_length`; the split count never exceeds 32 or the split count the grouped decode workspace was sized for. | +| | | | `N > 0` | `N` 64-candidate tiles per split (the sweep around the automatic value was flat; used by the tests to force many or few splits). | When `split_count > 32` (floating-point caches only) the generic `MergeDynamicSparseAttentionKernel` is used regardless of `ORT_DSA_SHARED_SCALE_MERGE`. The tests @@ -81,6 +94,8 @@ buffer is misaligned: | INT8 grouped decode vector staging (`ORT_DSA_INT8_VECTOR_STAGING=1`) | `key_scale`, `value_scale`, `auxiliary_key`, `auxiliary_value`: 16 bytes; `present_key`, `present_value`: 8 bytes | Scalar per-channel staging (`kVectorizedStaging=false`). | | Floating-point grouped decode and grouped prefill | Each staged K/V row: 16 bytes (checked per row in the kernel) | Element-wise copy into shared memory. | | Tensor-core BF16 prefill (`ORT_DSA_PREFILL_KERNEL=1`) | `present_key`, `present_value`: 16 bytes (8 for INT8); `key_scale`, `value_scale`: 16 bytes; `output`: 4 bytes | Generic fused/split or grouped prefill kernels. | +| Tensor-core INT8-cache prefill (`ORT_DSA_PREFILL_INT8_TC=1`) | In addition to the row above: `present_key`, `present_value`: 16 bytes; `output`: 8 bytes | `DynamicSparseAttentionPrefillKernel`. | +| Tensor-core INT8-cache decode (`ORT_DSA_DECODE_INT8_TC=1`) | `prepared_query`, `present_key`, `present_value`, `key_scale`, `value_scale`, attention workspace: 16 bytes | The CUDA-core grouped decode kernel. | Buffers allocated by ONNX Runtime are 256-byte aligned, so only IOBinding with offset buffers takes a fallback. The `*MisalignedBuffers_CUDA` tests cover each case. @@ -96,5 +111,26 @@ only as fallbacks: | `ORT_DSA_INT8_VECTOR_STAGING=0` | about +0.3 ms (+2%) slower | | `ORT_DSA_SHARED_SCALE_MERGE=0` | about +0.05 to +0.1 ms (+0.3 to 0.5%) slower | +Prefill with the INT8 KV cache (24 query / 2 KV heads, head size 256, 12 full-attention layers), H200: +`DynamicSparseAttentionInt8TcKernel` takes 2.15 ms per launch instead of 6.1 ms in a standalone +harness (8192 rows, 2048 candidates, 2.83x) and 2.8 ms instead of 7.9 ms per 8192-row launch in the +model. The previous kernel spent 2.8% of its instructions on `mma` and the rest on in-kernel +dequantization and index handling. Time to first token in the Qwen3.8-Flash-Next packages (one round) +drops by 4.0 to 4.8% at 2048, 8.4 to 9.5% at 8192 and 7.6 to 9.4% at 32768 prompt tokens; decode is +unchanged. The RMS error against an FP64 reference is 1.6e-3 (previous kernel 1.7e-3). +`ORT_DSA_PREFILL_INT8_TC=0` restores the previous kernel. The test +`DynamicSparseAttentionTest.Int8CachePrefillTensorCoreKernel_CUDA` compares both kernels with a +reference over head counts, candidate counts and invalid indices. + +Decode with the INT8 KV cache (`sequence_length <= 8`, 12 query heads per KV head, head size 256), H200: the +CUDA-core grouped decode kernel was instruction-issue and shared-load bound on its per-candidate serial chain, +70.3 us per layer plus 5.3 us for the merge (8 rows, 2048 selected candidates, flat from 2K to 32K context). The +split-K tensor-core kernel takes 9.5 us plus 3.5 us for the merge. In the Qwen3.8-Flash-Next packages the +speculative-decoding cycle (S = 8) is 4.6 to 6.8% shorter at 2K, 8K and 32K context. The tensor-core kernel +accumulates Q.K and P.V in a different FP32 order, so outputs are close but not bit-identical to the previous kernel; +`ORT_DSA_DECODE_INT8_TC=0` restores it exactly. The tests `Int8CacheDecodeTensorCoreKernel_CUDA` (reference +comparison over splits, candidate counts and invalid indices), `Int8DecodeQueryBeyondHalfRange_CUDA` (FP32 fallback +of the splits) and `DecodeSessionOptions_CUDA` (switch, tuning knob and invalid values) cover it. + The corresponding indexer switch results are in [../packed_sparse_attention_indexer.md](../packed_sparse_attention_indexer.md#benchmark). diff --git a/docs/contrib_ops/cuda/hyper_connection_low_rank_mix.md b/docs/contrib_ops/cuda/hyper_connection_low_rank_mix.md index 348131a3db9b3..8c0ec5331d088 100644 --- a/docs/contrib_ops/cuda/hyper_connection_low_rank_mix.md +++ b/docs/contrib_ops/cuda/hyper_connection_low_rank_mix.md @@ -37,6 +37,27 @@ Read once per process, so set it before starting it. weight inputs (`norm_scale`, `down_weight`, `up_weight`, `inject_weight`) are constant initializers; `0` forces plain launches. +`ProjectedSigmoidGateMulKernel` +([gated_add_impl.cu](../../../onnxruntime/contrib_ops/cuda/bert/gated_add_impl.cu)), +which precedes `HcDownKernel` in the Qwen3.8 decode graph, issues +`griddepcontrol.launch_dependents` at its start on `SM90+`. When `HcDownKernel` +is launched with PDL, its prologue (constant-weight prefetch) then overlaps with +the gate-mul kernel, while its reads of the gate-mul output still wait for +`griddepcontrol.wait`. The trigger has no effect for plain launches, changes no +result and has no switch. Together with the same early trigger in the fpA_intB GEMV +kernels (MatMulNBits) it shortens a Qwen3.8-Flash-Next `S = 8` decode cycle by 1.4-1.7%. + +`SigmoidF32` evaluates `expf(-|v|)` once and selects the numerator, +`e = expf(-|v|); (v >= 0 ? 1 : e) / (1 + e)`. The textbook two-branch form passes the +same argument to `expf` (`-v` for `v >= 0`, `v` otherwise), so the result is +bit-identical for all 2^32 `float` values (checked exhaustively in a standalone harness, +no fast-math). The branch itself was the cost: in the prefill pre-mix neighbouring +elements have different signs, so nearly every warp executed both sides. For hidden +2560, 4 branches and an 8192-token prefill chunk on an H200, `PreMixVec8Kernel` takes +82.7 µs instead of 112.1 µs per launch and `ScaledSiluKernel` 6.1 µs instead of 7.3 µs, +which shortens the TTFT by 0.3-1.1%. The fused decode kernels call the same function for +at most eight rows, and decode time is unchanged within the noise of 0.3%. + The rank-320 decode test (`HyperConnectionLowRankMixTest.FusedDecodeShapeRank320`) covers 1, 2, 3, 4, 7 and 8 rows. From 204d6593c0eb57f0a4fc9f8bcec19a3a2d2c9e62 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 29 Oct 2026 03:55:27 +0000 Subject: [PATCH 08/14] [CUDA] DynamicSparseAttention INT8 tensor-core kernels: exact for large softmax scales and tiny probabilities Address review feedback on the INT8 tensor-core prefill and split-K decode kernels. * Rounding Q * key_scale to fp16 flushes products below 2^-25, which moves a score by at most 2^-10 * scale. The kernels are now dispatched only for scale <= 1 (kInt8TcMaxScale); a larger scale uses the BF16-dequantizing prefill / CUDA-core decode kernels, which have no such limit. * The fp16 probability tile holds 2^15 * p instead of p, so a candidate keeps its contribution down to p = 2^-40 (it was lost below 2^-25) when a large value_scale makes it matter. The epilogue removes the factor from the prefill output and from the split-K partial output; the row sum stays unscaled fp32 and the merge kernel is unchanged. Probabilities of at least 2^-14 round exactly as before. * The session-option cases that expect a failure are not compiled in ORT_NO_EXCEPTIONS builds, where ORT_ENFORCE cannot be caught. New tests, each running prefill and decode: Int8TensorCoreLargeSoftmaxScale_CUDA (0.5 before, 1 / (1 + e^2) now) and Int8TensorCoreTinyProbabilityLargeValueScale_CUDA (0 before, 2^28 * e^-20 / (1 + e^-20) now). The docs list the new dispatch condition. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cuda/dynamic_sparse_attention.md | 6 +- .../bert/dynamic_sparse_attention_impl.cu | 11 +- .../dynamic_sparse_attention_prefill_int8.cuh | 41 +++++-- .../dynamic_sparse_attention_op_test.cc | 111 +++++++++++++++++- 4 files changed, 152 insertions(+), 17 deletions(-) diff --git a/docs/contrib_ops/cuda/dynamic_sparse_attention.md b/docs/contrib_ops/cuda/dynamic_sparse_attention.md index dad9e1ed4d7d6..7ed83c967d4c4 100644 --- a/docs/contrib_ops/cuda/dynamic_sparse_attention.md +++ b/docs/contrib_ops/cuda/dynamic_sparse_attention.md @@ -70,9 +70,9 @@ and 8 otherwise. | | | | `0` | `MergeGroupedDecodeDynamicSparseAttentionKernel`: per-channel `expf` over all splits. | | `ORT_DSA_PREFILL_KERNEL` | `1` | BF16, not grouped decode, `sequence_length >= 8`, `head_size == 256`, `selected_only` from the main cache, `smooth_softmax = 0`, SM80+ with enough opt-in shared memory | `1` | `DynamicSparseAttentionPrefillKernel` tensor-core prefill. | | | | | `0` | Generic fused/split or grouped prefill kernels. | -| `ORT_DSA_PREFILL_INT8_TC` | `1` | The `ORT_DSA_PREFILL_KERNEL` conditions with an INT8 main cache, 16-byte-aligned `present_key` and `present_value`, an 8-byte-aligned `output` and enough opt-in shared memory for `dsa_prefill::kInt8SharedBytes` | `1` | `DynamicSparseAttentionInt8TcKernel` ([dynamic_sparse_attention_prefill_int8.cuh](../../../onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh)): the gathered K/V rows stay INT8 in shared memory and are converted exactly to FP16 in registers for FP16 `mma` with FP32 accumulation. `key_scale` is folded into the query once per block and `value_scale` is applied to the output once. A block whose scaled query does not fit FP16 (`|Q * key_scale| > 65504`, inf or NaN) is evaluated by an FP32 path (`AttendRowFp32`) instead of being clipped, which would break the cancellation between large terms. That path is about 48x slower for the affected blocks (8192 rows on an H200: 2.75 ms to 131 ms); set `ep.cuda.dsa_prefill_int8_tc=0` for models whose scaled queries exceed the FP16 range, because the kernel below has no such limit. | +| `ORT_DSA_PREFILL_INT8_TC` | `1` | The `ORT_DSA_PREFILL_KERNEL` conditions with an INT8 main cache, 16-byte-aligned `present_key` and `present_value`, an 8-byte-aligned `output`, a softmax `scale <= 1` and enough opt-in shared memory for `dsa_prefill::kInt8SharedBytes` | `1` | `DynamicSparseAttentionInt8TcKernel` ([dynamic_sparse_attention_prefill_int8.cuh](../../../onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh)): the gathered K/V rows stay INT8 in shared memory and are converted exactly to FP16 in registers for FP16 `mma` with FP32 accumulation. `key_scale` is folded into the query once per block and `value_scale` is applied to the output once. Rounding `Q * key_scale` to FP16 flushes products below `2^-25`, which moves a score by at most `2^-10 * scale`, so a call with `scale > 1` (the default is `1 / sqrt(head_size)`) uses the kernel below. The FP16 probability tile holds `2^15 * p`, which the epilogue removes, so that a candidate keeps its contribution down to `p = 2^-40` (plain FP16 flushes below `2^-25`); the row sum stays unscaled FP32. A block whose scaled query does not fit FP16 (`abs(Q * key_scale) > 65504`, inf or NaN) is evaluated by an FP32 path (`AttendRowFp32`) instead of being clipped, which would break the cancellation between large terms. That path is about 48x slower for the affected blocks (8192 rows on an H200: 2.75 ms to 131 ms); set `ep.cuda.dsa_prefill_int8_tc=0` for models whose scaled queries exceed the FP16 range, because the kernel below has no such limit. | | | | | `0` | `DynamicSparseAttentionPrefillKernel`: dequantizes every gathered tile to BF16 in shared memory. | -| `ORT_DSA_DECODE_INT8_TC` | `1` | Grouped decode (conditions above) with BF16, an INT8 main cache, `head_size == 256`, `selected_only` from the main cache (no `local_plus_selected`), `smooth_softmax = 0`, at most `kMaxGroup` query heads per KV head, 16-byte-aligned `prepared_query`, `present_key`, `present_value`, `key_scale`, `value_scale` and attention workspace, SM80+ and enough opt-in shared memory for `dsa_prefill::kInt8SharedBytes` | `1` | `DynamicSparseAttentionInt8TcKernel`: the tensor-core pipeline of the row above, split over the 64-candidate tiles of each (row, KV head). Each CTA of the (row, KV head, split) grid writes an un-normalized partial output, the natural-log row maximum and the row sum in the layout of the grouped decode merge kernel, which is reused unchanged. The split count targets 384 CTAs (one wave of three per SM, at most 32 splits). A CTA whose scaled query does not fit FP16 evaluates its candidates in FP32 (`AttendRowFp32`) and publishes an FP32 partial in the same layout, so the range-safe fallback also holds for the splits. | +| `ORT_DSA_DECODE_INT8_TC` | `1` | Grouped decode (conditions above) with BF16, an INT8 main cache, `head_size == 256`, `selected_only` from the main cache (no `local_plus_selected`), `smooth_softmax = 0`, a softmax `scale <= 1`, at most `kMaxGroup` query heads per KV head, 16-byte-aligned `prepared_query`, `present_key`, `present_value`, `key_scale`, `value_scale` and attention workspace, SM80+ and enough opt-in shared memory for `dsa_prefill::kInt8SharedBytes` | `1` | `DynamicSparseAttentionInt8TcKernel`: the tensor-core pipeline of the row above, split over the 64-candidate tiles of each (row, KV head). Each CTA of the (row, KV head, split) grid writes an un-normalized partial output, the natural-log row maximum and the row sum in the layout of the grouped decode merge kernel, which is reused unchanged; the `2^15` probability factor of the row above is removed before a partial is written. The split count targets 384 CTAs (one wave of three per SM, at most 32 splits). A CTA whose scaled query does not fit FP16 evaluates its candidates in FP32 (`AttendRowFp32`) and publishes an FP32 partial in the same layout, so the range-safe fallback also holds for the splits. | | | | | `0` | The CUDA-core `GroupedDecodeDynamicSparseAttentionKernel`, bit for bit as before. | | `ORT_DSA_DECODE_INT8_TC_TILES_PER_SPLIT` | `0` | `ORT_DSA_DECODE_INT8_TC=1` and the conditions above | `0` | Automatic: `max(1, ceil(tiles * rows * kv_heads / 384))` tiles per split, with `rows = batch_size * sequence_length`; the split count never exceeds 32 or the split count the grouped decode workspace was sized for. | | | | | `N > 0` | `N` 64-candidate tiles per split (the sweep around the automatic value was flat; used by the tests to force many or few splits). | @@ -131,6 +131,8 @@ accumulates Q.K and P.V in a different FP32 order, so outputs are close but not `ORT_DSA_DECODE_INT8_TC=0` restores it exactly. The tests `Int8CacheDecodeTensorCoreKernel_CUDA` (reference comparison over splits, candidate counts and invalid indices), `Int8DecodeQueryBeyondHalfRange_CUDA` (FP32 fallback of the splits) and `DecodeSessionOptions_CUDA` (switch, tuning knob and invalid values) cover it. +`Int8TensorCoreLargeSoftmaxScale_CUDA` (the `scale > 1` fallback) and +`Int8TensorCoreTinyProbabilityLargeValueScale_CUDA` (the FP16 probability range) cover prefill and decode. The corresponding indexer switch results are in [../packed_sparse_attention_indexer.md](../packed_sparse_attention_indexer.md#benchmark). diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu index ce82f26c38877..be5c6dd398fba 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu @@ -2130,10 +2130,12 @@ Status LaunchDynamicSparseAttention( use_prefill_kernel = attributes.ptxVersion >= 80; } // The tensor-core INT8 variant converts the cached rows in registers; it needs 16-byte cache rows and 8-byte - // output vectors (the BF16-dequantizing kernel above is the fallback). + // output vectors, and a softmax scale of at most kInt8TcMaxScale so that it does not amplify its fp16 rounding + // of Q * key_scale (the BF16-dequantizing kernel above is the fallback). The op validates the scale as finite and + // non-negative. use_int8_tc_prefill = use_prefill_kernel && data.quantized_cache && - data.prefill_int8_tc && + data.prefill_int8_tc && parameters.scale <= dsa_prefill::kInt8TcMaxScale && IsAligned(data.present_key, 16) && IsAligned(data.present_value, 16) && IsAligned(data.output, 8) && static_cast(optin_shared) >= dsa_prefill::kInt8SharedBytes; if (use_int8_tc_prefill) { @@ -2157,8 +2159,11 @@ Status LaunchDynamicSparseAttention( constexpr int kMaxGridDimY = 65535; // The kernel loads 16-byte query / scale vectors, gathers 16-byte cache chunks and stores 16-byte partials; // buffers bound through IOBinding can start at any element offset, so misaligned ones use the generic kernel. + // Like the prefill variant it needs a softmax scale of at most kInt8TcMaxScale, which keeps the fp16 rounding of + // Q * key_scale from being amplified. use_int8_tc_decode = - cc_major >= 8 && parameters.head_size == dsa_prefill::kHeadDim && !local_plus_selected && + cc_major >= 8 && parameters.scale <= dsa_prefill::kInt8TcMaxScale && + parameters.head_size == dsa_prefill::kHeadDim && !local_plus_selected && parameters.selected_kv_source == DynamicSparseAttentionKvSource::kMain && !parameters.use_smooth_softmax && parameters.max_selected > 0 && parameters.kv_num_heads <= kMaxGridDimY && parameters.num_heads / parameters.kv_num_heads <= dsa_prefill::kMaxGroup && diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh index 4a518dc618619..dcdbbf0991e6b 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh @@ -13,13 +13,16 @@ // // The scales of the INT8 cache are per (kv head, channel) constants, so they leave the tile loop: // * the key scale is folded into the query once per CTA: Q' = Q * key_scale (rounded once to fp16), S = Q' K_int8; -// a CTA whose Q' does not fit fp16 (or is not finite) evaluates its query heads in fp32 instead (AttendRowFp32); +// a CTA whose Q' does not fit fp16 (or is not finite) evaluates its query heads in fp32 instead (AttendRowFp32). +// Products below 2^-25 flush to zero, which moves a score by at most 2^-10 * scale (256 channels, |K| <= 128), so +// the dispatcher selects this kernel only for scale <= kInt8TcMaxScale; // * the value scale is applied to the output once per CTA: O = (P V_int8) * value_scale. // The gathered int8 rows stay int8 in shared memory (cp.async, double buffered, one tile ahead). A fragment read with // ldmatrix is converted to fp16 exactly in registers (1024 + (byte ^ 0x80) - 1152, "magic number" conversion) and // consumed by fp16 MMAs with fp32 accumulation, so the per-tile dequantization pass, its shared memory round trip and // its barriers of the BF16-dequantizing kernel disappear. The scores stay in registers; the only cross-warp exchanges -// per tile are the row maxima and the fp16 probability tile. +// per tile are the row maxima and the fp16 probability tile, which holds kInt8ProbScale * p so that the tail of the +// softmax stays out of the fp16 subnormal range. // // The channel (reduction) dimension of S = Q' K^T is permuted so that one 4-byte ldmatrix element of a K row supplies // both B registers of an MMA k step: slot 2t+e <-> channel 16j + 4t + e and slot 2t+8+e <-> channel 16j + 4t + 2 + e, @@ -44,6 +47,19 @@ constexpr int kInt8QStride = kHeadDim + 8; // fp16 elements of the staged query constexpr int kInt8PStride = kTile + 8; // fp16 elements constexpr int kIdxStages = 3; +// The fp16 probability tile holds kInt8ProbScale * p for p in [0, 1]. Without the factor a candidate whose probability +// is below 2^-25 rounds to zero although a large value scale can make its contribution matter; with it the flush point +// drops to 2^-40 and the tile stays below the fp16 maximum. The PV accumulator carries the factor until the epilogue +// removes it; the softmax denominator is accumulated in fp32 without it. +constexpr float kInt8ProbScale = 32768.0f; // 2^15 +constexpr float kInt8ProbInverseScale = 1.0f / kInt8ProbScale; + +// Largest softmax scale served by the kernel below (the op validates it as finite and non-negative). Rounding +// Q * key_scale to fp16 flushes products below 2^-25 to zero, which moves a dot product by at most +// 256 * 128 * 2^-25 = 2^-10 before the scale multiplies it; a larger scale would amplify that error, so the dispatcher +// uses the BF16-dequantizing prefill / CUDA-core decode kernels for it. +constexpr float kInt8TcMaxScale = 1.0f; + // sK8[2] + sV8[2] (sQ aliases sK8[0] before the tile loop) + candidate indices + row maxima / sums + probabilities. constexpr size_t kInt8SharedBytes = 4 * static_cast(kInt8TileBytes) + kIdxStages * kTile * sizeof(int32_t) + kMaxGroup * 4 * sizeof(float) + kMaxGroup * kInt8PStride * sizeof(__half); @@ -266,7 +282,8 @@ __global__ void __launch_bounds__(kThreads, 3) DynamicSparseAttentionInt8TcKerne // Q' = Q * key_scale as fp16; the padding rows are zero so their (ignored) scores stay finite. A row whose Q' does // not fit fp16 (including inf / NaN) is evaluated in fp32 by the whole CTA instead. Underflow to fp16 subnormals - // costs at most 2^-25 per element, below the rounding of the fp16 operands, so it stays on this path. + // changes a product by at most 2^-25, i.e. a score by at most 2^-10 * scale, which the dispatcher keeps at the level + // of the fp16 rounding by requiring scale <= kInt8TcMaxScale. const __nv_bfloat16* query_group = query + (static_cast(row) * num_heads + kv_head * group) * kHeadDim; bool needs_fp32 = false; { @@ -480,10 +497,10 @@ __global__ void __launch_bounds__(kThreads, 3) DynamicSparseAttentionInt8TcKerne run_max0 = new_max0; run_max1 = new_max1; { - const __half2 h00 = __floats2half2_rn(p00, p01); - const __half2 h01 = __floats2half2_rn(p02, p03); - const __half2 h10 = __floats2half2_rn(p10, p11); - const __half2 h11 = __floats2half2_rn(p12, p13); + const __half2 h00 = __floats2half2_rn(p00 * kInt8ProbScale, p01 * kInt8ProbScale); + const __half2 h01 = __floats2half2_rn(p02 * kInt8ProbScale, p03 * kInt8ProbScale); + const __half2 h10 = __floats2half2_rn(p10 * kInt8ProbScale, p11 * kInt8ProbScale); + const __half2 h11 = __floats2half2_rn(p12 * kInt8ProbScale, p13 * kInt8ProbScale); *reinterpret_cast<__half2*>(sP + r0 * kInt8PStride + cbase) = h00; *reinterpret_cast<__half2*>(sP + r0 * kInt8PStride + cbase + 8) = h01; *reinterpret_cast<__half2*>(sP + (r0 + 8) * kInt8PStride + cbase) = h10; @@ -555,8 +572,9 @@ __global__ void __launch_bounds__(kThreads, 3) DynamicSparseAttentionInt8TcKerne const float4 ls1 = *reinterpret_cast(sM + (r0 + 8) * 4); const float sum0 = (ls0.x + ls0.y) + (ls0.z + ls0.w); const float sum1 = (ls1.x + ls1.y) + (ls1.z + ls1.w); - const float inv0 = sum0 == 0.0f ? 0.0f : 1.0f / sum0; - const float inv1 = sum1 == 0.0f ? 0.0f : 1.0f / sum1; + // acc_o carries kInt8ProbScale, which the normalization removes. + const float inv0 = sum0 == 0.0f ? 0.0f : kInt8ProbInverseScale / sum0; + const float inv1 = sum1 == 0.0f ? 0.0f : kInt8ProbInverseScale / sum1; if constexpr (kSplitK) { // Un-normalized partial of this split in the natural logarithm domain, as the merge kernel expects. Thread t owns @@ -578,7 +596,10 @@ __global__ void __launch_bounds__(kThreads, 3) DynamicSparseAttentionInt8TcKerne #pragma unroll for (int chunk = 0; chunk < 4; ++chunk) { const int dim = warp * 64 + chunk * 16 + t * 4; - const float4 vs = *reinterpret_cast(value_scale + kv_head * kHeadDim + dim); + // acc_o carries kInt8ProbScale (the partial_sum above does not), which the value scale takes out. + const float4 vs_raw = *reinterpret_cast(value_scale + kv_head * kHeadDim + dim); + const float4 vs = make_float4(vs_raw.x * kInt8ProbInverseScale, vs_raw.y * kInt8ProbInverseScale, + vs_raw.z * kInt8ProbInverseScale, vs_raw.w * kInt8ProbInverseScale); if (r0 < group) { *reinterpret_cast(partial_output + static_cast(partial_block0) * kHeadDim + dim) = make_float4(acc_o[2 * chunk][0] * vs.x, acc_o[2 * chunk + 1][0] * vs.y, acc_o[2 * chunk][1] * vs.z, diff --git a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc index 104d0582391c4..82cd682f4d949 100644 --- a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc @@ -811,14 +811,110 @@ TEST(DynamicSparseAttentionTest, Int8PrefillQueryBeyondHalfRange_CUDA) { OpTester::ExpectResult::kExpectSuccess, "", configs); } +// Two selected candidates (cache positions 0 and 1) of an otherwise empty INT8 cache, so that the scores and the output +// follow from a few numbers: every query head holds `query` on channel 0, the keys hold key_codes on channel 0 and the +// values hold value_codes on every channel. The first row of the case selects nothing and outputs zeros. +struct TwoCandidateInt8Spec { + float query; + float key_scale; + float value_scale; + float scale; + float key_codes[2]; + float value_codes[2]; +}; + +DynamicSparseAttentionCase MakeTwoCandidateInt8Case(const TwoCandidateInt8Spec& spec, int num_heads, + int sequence_length) { + auto c = MakeSelectedOnlyPrefillCase(true, num_heads, 2, false, sequence_length); + c.scale = spec.scale; + std::fill(c.query.begin(), c.query.end(), 0.0f); + std::fill(c.key_scale.begin(), c.key_scale.end(), spec.key_scale); + std::fill(c.value_scale.begin(), c.value_scale.end(), spec.value_scale); + std::fill(c.key.begin(), c.key.end(), 0.0f); + std::fill(c.value.begin(), c.value.end(), 0.0f); + std::fill(c.past_key.begin(), c.past_key.end(), 0.0f); + std::fill(c.past_value.begin(), c.past_value.end(), 0.0f); + for (int b = 0; b < c.batch_size; ++b) { + for (int kv = 0; kv < c.kv_num_heads; ++kv) { + const int64_t offset = (b * c.kv_num_heads + kv) * c.cache_sequence_length * c.head_size; + for (int i = 0; i < 2; ++i) { + c.past_key[offset + i * c.head_size] = spec.key_codes[i]; + std::fill_n(c.past_value.begin() + offset + i * c.head_size, c.head_size, spec.value_codes[i]); + } + } + } + c.expected_present_key = c.past_key; + c.expected_present_value = c.past_value; + for (size_t i = 0; i < c.query.size(); i += c.head_size) c.query[i] = spec.query; + ComputeSelectedOnlyPrefillExpectedOutput(c, true); + return c; +} + +// Every row but the first outputs `expected` on every channel. Prefill (9 rows per batch) and decode (3 rows) run +// with the tensor-core INT8 kernels enabled. +void RunTwoCandidateInt8Case(const TwoCandidateInt8Spec& spec, float expected) { + for (bool decode : {false, true}) { + SCOPED_TRACE(decode ? "decode" : "prefill"); + const auto c = MakeTwoCandidateInt8Case(spec, 12, decode ? 3 : 9); + for (size_t i = 0; i < c.expected_output.size(); i += c.head_size) { + const size_t row = i / (c.num_heads * c.head_size); + ASSERT_NEAR(c.expected_output[i], row == 0 ? 0.0f : expected, 1e-6f); + } + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), + OpTester::ExpectResult::kExpectSuccess, "", + decode ? DecodeKernelConfigs("1") : PrefillKernelConfigs("1", "1")); + if (::testing::Test::HasFailure()) return; + } +} + +// The tensor-core INT8 kernels round Q * key_scale to fp16, which flushes products below 2^-25 to zero; a softmax scale +// above 1 would amplify that error, so those calls use the BF16-dequantizing prefill and CUDA-core decode kernels. +// Every product here is 2^-30 (zero in fp16) and the scale of 2^30 turns it into scores of +-1, so the output is +// 1 / (1 + e^2) rather than the 0.5 an fp16 flush gives. +TEST(DynamicSparseAttentionTest, Int8TensorCoreLargeSoftmaxScale_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + TwoCandidateInt8Spec spec{}; + spec.query = 1.0f; + spec.key_scale = 0x1p-30f; + spec.value_scale = 1.0f; + spec.scale = 0x1p30f; + spec.key_codes[0] = 1.0f; + spec.key_codes[1] = -1.0f; + spec.value_codes[1] = 1.0f; + RunTwoCandidateInt8Case(spec, static_cast(1.0 / (1.0 + std::exp(2.0)))); +} + +// The tensor-core INT8 kernels pass the softmax probabilities to the PV mma as fp16 holding 2^15 * p (the factor leaves +// in the epilogue), so a candidate keeps its contribution down to p = 2^-40 where plain fp16 flushes p below 2^-25. The +// second candidate has probability e^-20 (score 320 * -1 / 16) and is the only one with a non-zero value, whose scale is +// 2^28: the output is 2^28 * e^-20 / (1 + e^-20), not zero. The expected value is below 1, where the BF16 spacing is +// smaller than the output tolerance. +TEST(DynamicSparseAttentionTest, Int8TensorCoreTinyProbabilityLargeValueScale_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + TwoCandidateInt8Spec spec{}; + spec.query = 320.0f; + spec.key_scale = 1.0f; + spec.value_scale = 0x1p28f; + spec.scale = 0.0625f; + spec.key_codes[1] = -1.0f; + spec.value_codes[1] = 1.0f; + const double tail = std::exp(-20.0); + RunTwoCandidateInt8Case(spec, static_cast(0x1p28 * tail / (1.0 + tail))); +} + // Explicit ep.cuda.* entries take precedence over the environment, which is consulted only when no entry is set; an -// invalid value from either source fails kernel construction. +// invalid value from either source fails kernel construction (with an exception, so those cases are skipped in builds +// without exceptions). TEST(DynamicSparseAttentionTest, PrefillSessionOptions_CUDA) { if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; const auto c = MakeSelectedOnlyPrefillCase(true, 12, 65, true); const auto expect_success = OpTester::ExpectResult::kExpectSuccess; +#ifndef ORT_NO_EXCEPTIONS const auto expect_failure = OpTester::ExpectResult::kExpectFailure; +#endif { ScopedEnvironmentVariables env{ @@ -830,23 +926,27 @@ TEST(DynamicSparseAttentionTest, PrefillSessionOptions_CUDA) { PrefillKernelConfigs(kernel, int8_tc)); } } +#ifndef ORT_NO_EXCEPTIONS RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_failure, "ORT_DSA_PREFILL_INT8_TC", {{kOrtSessionOptionsCudaDsaPrefillKernel, "1"}}); RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_failure, "ORT_DSA_PREFILL_KERNEL", {{kOrtSessionOptionsCudaDsaPrefillInt8Tc, "1"}}); +#endif } for (const char* int8_tc : {"0", "1"}) { ScopedEnvironmentVariables env{EnvVarMap{{"ORT_DSA_PREFILL_KERNEL", "1"}, {"ORT_DSA_PREFILL_INT8_TC", int8_tc}}}; RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_success); } +#ifndef ORT_NO_EXCEPTIONS for (const char* invalid : {"2", "-1", "", "1x", "true"}) { SCOPED_TRACE(std::string("value='") + invalid + "'"); RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_failure, kOrtSessionOptionsCudaDsaPrefillInt8Tc, {{kOrtSessionOptionsCudaDsaPrefillInt8Tc, invalid}}); } +#endif } // The split-K tensor-core INT8 decode kernel (ep.cuda.dsa_decode_int8_tc=1) and the CUDA-core grouped decode kernel @@ -937,13 +1037,16 @@ TEST(DynamicSparseAttentionTest, Int8DecodeQueryBeyondHalfRange_CUDA) { } // Explicit ep.cuda.* entries take precedence over the environment, which is consulted only when no entry is set; an -// invalid value from either source fails kernel construction. +// invalid value from either source fails kernel construction (with an exception, so those cases are skipped in builds +// without exceptions). TEST(DynamicSparseAttentionTest, DecodeSessionOptions_CUDA) { if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; const auto c = MakeSelectedOnlyPrefillCase(true, 12, 65, true, 3); const auto expect_success = OpTester::ExpectResult::kExpectSuccess; +#ifndef ORT_NO_EXCEPTIONS const auto expect_failure = OpTester::ExpectResult::kExpectFailure; +#endif { ScopedEnvironmentVariables env{ @@ -955,18 +1058,21 @@ TEST(DynamicSparseAttentionTest, DecodeSessionOptions_CUDA) { DecodeKernelConfigs(int8_tc, tiles_per_split)); } } +#ifndef ORT_NO_EXCEPTIONS RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_failure, "(ORT_DSA_DECODE_INT8_TC_TILES_PER_SPLIT)", {{kOrtSessionOptionsCudaDsaDecodeInt8Tc, "1"}}); RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_failure, "(ORT_DSA_DECODE_INT8_TC)", {{kOrtSessionOptionsCudaDsaDecodeInt8TcTilesPerSplit, "0"}}); +#endif } for (const char* int8_tc : {"0", "1"}) { ScopedEnvironmentVariables env{ EnvVarMap{{"ORT_DSA_DECODE_INT8_TC", int8_tc}, {"ORT_DSA_DECODE_INT8_TC_TILES_PER_SPLIT", "2"}}}; RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_success); } +#ifndef ORT_NO_EXCEPTIONS for (const char* invalid : {"2", "-1", "", "1x", "true"}) { SCOPED_TRACE(std::string("value='") + invalid + "'"); RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_failure, @@ -979,6 +1085,7 @@ TEST(DynamicSparseAttentionTest, DecodeSessionOptions_CUDA) { kOrtSessionOptionsCudaDsaDecodeInt8TcTilesPerSplit, {{kOrtSessionOptionsCudaDsaDecodeInt8TcTilesPerSplit, invalid}}); } +#endif } TEST(DynamicSparseAttentionTest, FloatingCacheGroupedSingleSplitWorkspace_CUDA) { From ceed60271c8ce5d3feb5bcfb62f6682d08caf28f Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Fri, 30 Oct 2026 22:03:13 +0000 Subject: [PATCH 09/14] Gate INT8 TC attention on value_scale <= 2^16 with FP32 fallback When any channel of value_scale exceeds kInt8TcMaxValueScale (2^16), route the CTA through AttendRowFp32 (for both prefill and split-K decode) and fold value_scale into the FP32 numerator with an exponent shift for extreme negative log-probabilities so tiny probabilities amplified by large finite value scales are preserved. Extend Int8TensorCoreTinyProbabilityLargeValueScale_CUDA to cover TC and reference paths for 2^16, 2^28, and 2^60 value scales across prefill and split-K decode. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cuda/dynamic_sparse_attention.md | 2 +- .../dynamic_sparse_attention_prefill_int8.cuh | 52 +++++++++---- .../dynamic_sparse_attention_op_test.cc | 75 ++++++++++++++----- 3 files changed, 92 insertions(+), 37 deletions(-) diff --git a/docs/contrib_ops/cuda/dynamic_sparse_attention.md b/docs/contrib_ops/cuda/dynamic_sparse_attention.md index 7ed83c967d4c4..ed96cf3c727a7 100644 --- a/docs/contrib_ops/cuda/dynamic_sparse_attention.md +++ b/docs/contrib_ops/cuda/dynamic_sparse_attention.md @@ -70,7 +70,7 @@ and 8 otherwise. | | | | `0` | `MergeGroupedDecodeDynamicSparseAttentionKernel`: per-channel `expf` over all splits. | | `ORT_DSA_PREFILL_KERNEL` | `1` | BF16, not grouped decode, `sequence_length >= 8`, `head_size == 256`, `selected_only` from the main cache, `smooth_softmax = 0`, SM80+ with enough opt-in shared memory | `1` | `DynamicSparseAttentionPrefillKernel` tensor-core prefill. | | | | | `0` | Generic fused/split or grouped prefill kernels. | -| `ORT_DSA_PREFILL_INT8_TC` | `1` | The `ORT_DSA_PREFILL_KERNEL` conditions with an INT8 main cache, 16-byte-aligned `present_key` and `present_value`, an 8-byte-aligned `output`, a softmax `scale <= 1` and enough opt-in shared memory for `dsa_prefill::kInt8SharedBytes` | `1` | `DynamicSparseAttentionInt8TcKernel` ([dynamic_sparse_attention_prefill_int8.cuh](../../../onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh)): the gathered K/V rows stay INT8 in shared memory and are converted exactly to FP16 in registers for FP16 `mma` with FP32 accumulation. `key_scale` is folded into the query once per block and `value_scale` is applied to the output once. Rounding `Q * key_scale` to FP16 flushes products below `2^-25`, which moves a score by at most `2^-10 * scale`, so a call with `scale > 1` (the default is `1 / sqrt(head_size)`) uses the kernel below. The FP16 probability tile holds `2^15 * p`, which the epilogue removes, so that a candidate keeps its contribution down to `p = 2^-40` (plain FP16 flushes below `2^-25`); the row sum stays unscaled FP32. A block whose scaled query does not fit FP16 (`abs(Q * key_scale) > 65504`, inf or NaN) is evaluated by an FP32 path (`AttendRowFp32`) instead of being clipped, which would break the cancellation between large terms. That path is about 48x slower for the affected blocks (8192 rows on an H200: 2.75 ms to 131 ms); set `ep.cuda.dsa_prefill_int8_tc=0` for models whose scaled queries exceed the FP16 range, because the kernel below has no such limit. | +| `ORT_DSA_PREFILL_INT8_TC` | `1` | The `ORT_DSA_PREFILL_KERNEL` conditions with an INT8 main cache, 16-byte-aligned `present_key` and `present_value`, an 8-byte-aligned `output`, a softmax `scale <= 1` and enough opt-in shared memory for `dsa_prefill::kInt8SharedBytes` | `1` | `DynamicSparseAttentionInt8TcKernel` ([dynamic_sparse_attention_prefill_int8.cuh](../../../onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh)): the gathered K/V rows stay INT8 in shared memory and are converted exactly to FP16 in registers for FP16 `mma` with FP32 accumulation. `key_scale` is folded into the query once per block and `value_scale` is applied to the output once. Rounding `Q * key_scale` to FP16 flushes products below `2^-25`, which moves a score by at most `2^-10 * scale`, so a call with `scale > 1` (the default is `1 / sqrt(head_size)`) uses the kernel below. The FP16 probability tile holds `2^15 * p`, which the epilogue removes, so that a candidate keeps its contribution down to `p = 2^-40` (plain FP16 flushes below `2^-25`); the row sum stays unscaled FP32. A block whose scaled query does not fit FP16 (`abs(Q * key_scale) > 65504`, inf or NaN) or whose value scale exceeds `kInt8TcMaxValueScale` (`2^16`, above which probabilities flushed below `2^-40` can become significant after `value_scale`) is evaluated by an FP32 path (`AttendRowFp32`) instead of being clipped, which would break the cancellation between large terms. That path is about 48x slower for the affected blocks (8192 rows on an H200: 2.75 ms to 131 ms); set `ep.cuda.dsa_prefill_int8_tc=0` for models whose scaled queries exceed the FP16 range, because the kernel below has no such limit. | | | | | `0` | `DynamicSparseAttentionPrefillKernel`: dequantizes every gathered tile to BF16 in shared memory. | | `ORT_DSA_DECODE_INT8_TC` | `1` | Grouped decode (conditions above) with BF16, an INT8 main cache, `head_size == 256`, `selected_only` from the main cache (no `local_plus_selected`), `smooth_softmax = 0`, a softmax `scale <= 1`, at most `kMaxGroup` query heads per KV head, 16-byte-aligned `prepared_query`, `present_key`, `present_value`, `key_scale`, `value_scale` and attention workspace, SM80+ and enough opt-in shared memory for `dsa_prefill::kInt8SharedBytes` | `1` | `DynamicSparseAttentionInt8TcKernel`: the tensor-core pipeline of the row above, split over the 64-candidate tiles of each (row, KV head). Each CTA of the (row, KV head, split) grid writes an un-normalized partial output, the natural-log row maximum and the row sum in the layout of the grouped decode merge kernel, which is reused unchanged; the `2^15` probability factor of the row above is removed before a partial is written. The split count targets 384 CTAs (one wave of three per SM, at most 32 splits). A CTA whose scaled query does not fit FP16 evaluates its candidates in FP32 (`AttendRowFp32`) and publishes an FP32 partial in the same layout, so the range-safe fallback also holds for the splits. | | | | | `0` | The CUDA-core `GroupedDecodeDynamicSparseAttentionKernel`, bit for bit as before. | diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh index dcdbbf0991e6b..1cda8401d6471 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh @@ -16,7 +16,9 @@ // a CTA whose Q' does not fit fp16 (or is not finite) evaluates its query heads in fp32 instead (AttendRowFp32). // Products below 2^-25 flush to zero, which moves a score by at most 2^-10 * scale (256 channels, |K| <= 128), so // the dispatcher selects this kernel only for scale <= kInt8TcMaxScale; -// * the value scale is applied to the output once per CTA: O = (P V_int8) * value_scale. +// * the value scale is applied to the output once per CTA: O = (P V_int8) * value_scale; a CTA whose value_scale +// exceeds kInt8TcMaxValueScale (so that a probability flushed below 2^-40 by the fp16 probability tile could +// still move the scaled output) is also evaluated in fp32 by AttendRowFp32. // The gathered int8 rows stay int8 in shared memory (cp.async, double buffered, one tile ahead). A fragment read with // ldmatrix is converted to fp16 exactly in registers (1024 + (byte ^ 0x80) - 1152, "magic number" conversion) and // consumed by fp16 MMAs with fp32 accumulation, so the per-tile dequantization pass, its shared memory round trip and @@ -50,9 +52,12 @@ constexpr int kIdxStages = 3; // The fp16 probability tile holds kInt8ProbScale * p for p in [0, 1]. Without the factor a candidate whose probability // is below 2^-25 rounds to zero although a large value scale can make its contribution matter; with it the flush point // drops to 2^-40 and the tile stays below the fp16 maximum. The PV accumulator carries the factor until the epilogue -// removes it; the softmax denominator is accumulated in fp32 without it. +// removes it; the softmax denominator is accumulated in fp32 without it. When any channel has +// |value_scale| > kInt8TcMaxValueScale (so that |p * V * value_scale| for p <= 2^-40 can exceed 2^-17), the CTA takes +// AttendRowFp32 instead. constexpr float kInt8ProbScale = 32768.0f; // 2^15 constexpr float kInt8ProbInverseScale = 1.0f / kInt8ProbScale; +constexpr float kInt8TcMaxValueScale = 65536.0f; // 2^16 // Largest softmax scale served by the kernel below (the op validates it as finite and non-negative). Rounding // Q * key_scale to fp16 flushes products below 2^-25 to zero, which moves a dot product by at most @@ -125,12 +130,14 @@ __device__ __forceinline__ void Int8x4ToHalf2Strided(uint32_t w, uint32_t& even, } // FP32 evaluation of the query heads of one (row, KV head), one warp per head and eight channels per lane, for rows -// whose Q * key_scale leaves the FP16 range of the tensor-core path. Clipping there would break the cancellation -// between large terms. `query` and `output` point at the first head of the group; `key_scale` and `value_scale` at the -// KV head. The candidates [begin, end) of `row_indices` are consumed: all of them for prefill, the range of the CTA -// for a decode split. With kSplitK the un-normalized partial (running maximum in the natural logarithm domain, running -// sum, value-scaled accumulator) of block partial_block + head * split_count goes to partial_*, as the main path does; -// otherwise the normalized result goes to output. +// whose Q * key_scale leaves the FP16 range of the tensor-core path or whose value_scale exceeds kInt8TcMaxValueScale. +// Clipping Q * key_scale would break the cancellation between large terms, and a larger value_scale can amplify +// probabilities below the 2^-40 FP16 probability-tile flush boundary. `query` and `output` point at the first head of +// the group; `key_scale` and `value_scale` at the KV head. The candidates [begin, end) of `row_indices` are consumed: +// all of them for prefill, the range of the CTA for a decode split. With kSplitK the un-normalized partial (running +// maximum in the natural logarithm domain, running sum, value-scaled accumulator) of block +// partial_block + head * split_count goes to partial_*, as the main path does; otherwise the normalized result goes to +// output. template __device__ __noinline__ void AttendRowFp32(const __nv_bfloat16* __restrict__ query, const int8_t* __restrict__ key, const int8_t* __restrict__ value, const float* __restrict__ key_scale, @@ -141,9 +148,11 @@ __device__ __noinline__ void AttendRowFp32(const __nv_bfloat16* __restrict__ que float* __restrict__ partial_sum, float* __restrict__ partial_output, int partial_block, int split_count, int warp, int lane) { float ks[8]; + float vs[8]; #pragma unroll for (int j = 0; j < 8; ++j) { ks[j] = key_scale[lane + j * 32]; + vs[j] = value_scale[lane + j * 32]; } for (int h = warp; h < group; h += kThreads / 32) { float q[8]; @@ -172,11 +181,17 @@ __device__ __noinline__ void AttendRowFp32(const __nv_bfloat16* __restrict__ que const float score = dot * scale; const float next_max = fmaxf(maximum, score); const float correction = expf(maximum - next_max); - const float probability = expf(score - next_max); - denominator = denominator * correction + probability; + const float diff = score - next_max; + // Shift by 2^80 before expf when diff < -80 so tiny probabilities that value_scale brings back into range + // do not underflow fp32 before multiplying by vs[j]. + const bool shift_exp = diff < -80.0f; + const float probability = expf(shift_exp ? diff + 80.0f * 0.693147180559945309f : diff); + const float prob_unscale = shift_exp ? 0x1p-80f : 1.0f; + denominator = denominator * correction + (shift_exp ? 0.0f : probability); #pragma unroll for (int j = 0; j < 8; ++j) { - numerator[j] = numerator[j] * correction + probability * static_cast(value[offset + j * 32]); + const float pv = probability * static_cast(value[offset + j * 32]); + numerator[j] = numerator[j] * correction + pv * (vs[j] * prob_unscale); } maximum = next_max; } @@ -189,14 +204,14 @@ __device__ __noinline__ void AttendRowFp32(const __nv_bfloat16* __restrict__ que #pragma unroll for (int j = 0; j < 8; ++j) { const int d = lane + j * 32; - partial_output[static_cast(block) * kHeadDim + d] = numerator[j] * value_scale[d]; + partial_output[static_cast(block) * kHeadDim + d] = numerator[j]; } } else { const float inverse = denominator == 0.0f ? 0.0f : 1.0f / denominator; #pragma unroll for (int j = 0; j < 8; ++j) { const int d = lane + j * 32; - output[static_cast(h) * kHeadDim + d] = __float2bfloat16_rn(numerator[j] * inverse * value_scale[d]); + output[static_cast(h) * kHeadDim + d] = __float2bfloat16_rn(numerator[j] * inverse); } } } @@ -281,15 +296,20 @@ __global__ void __launch_bounds__(kThreads, 3) DynamicSparseAttentionInt8TcKerne } // Q' = Q * key_scale as fp16; the padding rows are zero so their (ignored) scores stay finite. A row whose Q' does - // not fit fp16 (including inf / NaN) is evaluated in fp32 by the whole CTA instead. Underflow to fp16 subnormals - // changes a product by at most 2^-25, i.e. a score by at most 2^-10 * scale, which the dispatcher keeps at the level - // of the fp16 rounding by requiring scale <= kInt8TcMaxScale. + // not fit fp16 (including inf / NaN) or whose value_scale exceeds kInt8TcMaxValueScale is evaluated in fp32 by the + // whole CTA instead. Underflow to fp16 subnormals changes a product by at most 2^-25, i.e. a score by at most + // 2^-10 * scale, which the dispatcher keeps at the level of the fp16 rounding by requiring scale <= kInt8TcMaxScale. const __nv_bfloat16* query_group = query + (static_cast(row) * num_heads + kv_head * group) * kHeadDim; bool needs_fp32 = false; { const int chunk = tid & 31; const float4 k0 = *reinterpret_cast(key_scale + kv_head * kHeadDim + chunk * 8); const float4 k1 = *reinterpret_cast(key_scale + kv_head * kHeadDim + chunk * 8 + 4); + if (tid < kHeadDim / 4) { + const float4 v = *reinterpret_cast(value_scale + kv_head * kHeadDim + tid * 4); + needs_fp32 = !(fabsf(v.x) <= kInt8TcMaxValueScale) || !(fabsf(v.y) <= kInt8TcMaxValueScale) || + !(fabsf(v.z) <= kInt8TcMaxValueScale) || !(fabsf(v.w) <= kInt8TcMaxValueScale); + } auto to_half2 = [&needs_fp32](float x, float y) { needs_fp32 |= !(fabsf(x) <= 65504.0f) || !(fabsf(y) <= 65504.0f); return __floats2half2_rn(x, y); diff --git a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc index 82cd682f4d949..a80c9440cc2ff 100644 --- a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc @@ -851,19 +851,21 @@ DynamicSparseAttentionCase MakeTwoCandidateInt8Case(const TwoCandidateInt8Spec& } // Every row but the first outputs `expected` on every channel. Prefill (9 rows per batch) and decode (3 rows) run -// with the tensor-core INT8 kernels enabled. +// with the tensor-core INT8 kernels both disabled (reference fallback) and enabled. void RunTwoCandidateInt8Case(const TwoCandidateInt8Spec& spec, float expected) { for (bool decode : {false, true}) { - SCOPED_TRACE(decode ? "decode" : "prefill"); const auto c = MakeTwoCandidateInt8Case(spec, 12, decode ? 3 : 9); for (size_t i = 0; i < c.expected_output.size(); i += c.head_size) { const size_t row = i / (c.num_heads * c.head_size); - ASSERT_NEAR(c.expected_output[i], row == 0 ? 0.0f : expected, 1e-6f); + ASSERT_NEAR(c.expected_output[i], row == 0 ? 0.0f : expected, 1e-6f * std::max(1.0f, std::fabs(expected))); + } + for (const char* int8_tc : {"0", "1"}) { + SCOPED_TRACE(std::string(decode ? "decode" : "prefill") + ", int8_tc=" + int8_tc); + RunDynamicSparseAttentionCase( + c, DefaultCudaExecutionProvider(), OpTester::ExpectResult::kExpectSuccess, "", + decode ? DecodeKernelConfigs(int8_tc) : PrefillKernelConfigs("1", int8_tc)); + if (::testing::Test::HasFailure()) return; } - RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), - OpTester::ExpectResult::kExpectSuccess, "", - decode ? DecodeKernelConfigs("1") : PrefillKernelConfigs("1", "1")); - if (::testing::Test::HasFailure()) return; } } @@ -886,22 +888,55 @@ TEST(DynamicSparseAttentionTest, Int8TensorCoreLargeSoftmaxScale_CUDA) { } // The tensor-core INT8 kernels pass the softmax probabilities to the PV mma as fp16 holding 2^15 * p (the factor leaves -// in the epilogue), so a candidate keeps its contribution down to p = 2^-40 where plain fp16 flushes p below 2^-25. The -// second candidate has probability e^-20 (score 320 * -1 / 16) and is the only one with a non-zero value, whose scale is -// 2^28: the output is 2^28 * e^-20 / (1 + e^-20), not zero. The expected value is below 1, where the BF16 spacing is -// smaller than the output tolerance. +// in the epilogue), so a candidate keeps its contribution down to p = 2^-40 where plain fp16 flushes p below 2^-25; +// when value_scale > 2^16 (kInt8TcMaxValueScale), the CTA evaluates the row/split in fp32 (AttendRowFp32) so that +// probabilities below 2^-40 amplified by arbitrarily large finite value scales are preserved as well: +// 1. query = 288 -> score -18, p = e^-18 ~= 2^-25.97 (flushes in unscaled fp16, kept by 2^15 * p on the TC path with +// value_scale = 2^16 <= kInt8TcMaxValueScale and value code 100); +// 2. query = 320 -> score -20, p = e^-20 ~= 2^-28.85, value_scale = 2^28; +// 3. query = 480 -> score -30, p = e^-30 ~= 2^-43.28 (flushes even 2^15 * p in fp16), value_scale = 2^60, output +// 2^60 * e^-30 / (1 + e^-30) ~= 107886.05 via the AttendRowFp32 fallback. TEST(DynamicSparseAttentionTest, Int8TensorCoreTinyProbabilityLargeValueScale_CUDA) { if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; - TwoCandidateInt8Spec spec{}; - spec.query = 320.0f; - spec.key_scale = 1.0f; - spec.value_scale = 0x1p28f; - spec.scale = 0.0625f; - spec.key_codes[1] = -1.0f; - spec.value_codes[1] = 1.0f; - const double tail = std::exp(-20.0); - RunTwoCandidateInt8Case(spec, static_cast(0x1p28 * tail / (1.0 + tail))); + { + SCOPED_TRACE("tc_prob_scale_2p16"); + TwoCandidateInt8Spec spec{}; + spec.query = 288.0f; + spec.key_scale = 1.0f; + spec.value_scale = 0x1p16f; + spec.scale = 0.0625f; + spec.key_codes[1] = -1.0f; + spec.value_codes[1] = 100.0f; + const double tail = std::exp(-18.0); + RunTwoCandidateInt8Case(spec, static_cast(100.0 * 0x1p16 * tail / (1.0 + tail))); + if (::testing::Test::HasFailure()) return; + } + { + SCOPED_TRACE("fp32_fallback_2p28"); + TwoCandidateInt8Spec spec{}; + spec.query = 320.0f; + spec.key_scale = 1.0f; + spec.value_scale = 0x1p28f; + spec.scale = 0.0625f; + spec.key_codes[1] = -1.0f; + spec.value_codes[1] = 1.0f; + const double tail = std::exp(-20.0); + RunTwoCandidateInt8Case(spec, static_cast(0x1p28 * tail / (1.0 + tail))); + if (::testing::Test::HasFailure()) return; + } + { + SCOPED_TRACE("fp32_fallback_2p60"); + TwoCandidateInt8Spec spec{}; + spec.query = 480.0f; + spec.key_scale = 1.0f; + spec.value_scale = 0x1p60f; + spec.scale = 0.0625f; + spec.key_codes[1] = -1.0f; + spec.value_codes[1] = 1.0f; + const double tail = std::exp(-30.0); + RunTwoCandidateInt8Case(spec, static_cast(0x1p60 * tail / (1.0 + tail))); + } } // Explicit ep.cuda.* entries take precedence over the environment, which is consulted only when no entry is set; an From 1bf2eb17cb2196fd27abf1a66c4fdeaa827afc2a Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 31 Oct 2026 20:38:02 +0000 Subject: [PATCH 10/14] [CUDA] DynamicSparseAttention INT8 tensor-core kernels: review fixes for numerics, split count and rollout Numerics of DynamicSparseAttentionInt8TcKernel (prefill and split-K decode): - Scores and the running maximum are kept in the natural-log domain and only the difference to the new maximum is scaled by log2(e) for ex2, so two large, nearby logits are not rounded apart by a base-2 conversion. - Q' = fp16(Q * key_scale) is rounded once. A CTA bounds the rounding with eps = scale * 128 * sum_d |Q'_d - fp16(Q'_d)| per query head and evaluates its heads in FP32 (AttendRowFp32) when eps exceeds kInt8TcQueryErrorBudget (0.125), besides the existing cases of Q' leaving the FP16 range or being non-finite. A query whose large terms cancel to a small logit exceeds the budget, an ordinary query does not (RC model: eps about 0.015). - Decode keeps the probability tile as an FP16 high part plus an FP16 residual (about 22 significant bits), because the signed sum of large value codes can cancel far below the rounding of one FP16 probability. Prefill keeps a single tile, like the BF16-dequantizing kernel. - A large finite value_scale (> 2^16) no longer sends every CTA to the FP32 path (one 8192-row launch took 136 ms instead of 2.8 ms). A CTA takes that path only when a valid candidate is more than 20 nats (kInt8TcTailDepth) below the running maximum, checked after the tile loop. A non-finite value_scale still takes it. Dispatch: - The split count of the decode kernel follows the CTAs that can be resident (SM count times cudaOccupancyMaxActiveBlocksPerMultiprocessor): the smallest power-of-two number of tiles per split with which all CTAs fit one wave, clamped to [1, 32] splits and to the workspace. This replaces the fixed target of 384 CTAs, which assumed 132 SMs with three CTAs each. A positive ep.cuda.dsa_decode_int8_tc_tiles_per_split replaces the automatic value. - Both INT8 tensor-core kernels are opt-in (ep.cuda.dsa_prefill_int8_tc and ep.cuda.dsa_decode_int8_tc, or ORT_DSA_PREFILL_INT8_TC and ORT_DSA_DECODE_INT8_TC, default 0), so the default dispatch and its results stay those of main until the final revision is qualified end to end. Options: the 0/1 switches, ORT_DSA_PREFILL_KERNEL included, accept exactly "0" and "1" ("00", "01", "-0" and an empty session-option string fail kernel construction). The tiles-per-split knob stays a non-negative decimal integer. The public key descriptions state the opt-in defaults and that the tiles-per-split value is a target. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../onnxruntime_session_options_config_keys.h | 16 +- .../cuda/bert/dynamic_sparse_attention.cc | 12 +- .../bert/dynamic_sparse_attention_impl.cu | 52 ++-- .../cuda/bert/dynamic_sparse_attention_impl.h | 4 +- .../dynamic_sparse_attention_prefill_int8.cuh | 262 +++++++++++++----- .../contrib_ops/cuda/cuda_session_options.h | 25 +- 6 files changed, 264 insertions(+), 107 deletions(-) diff --git a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h index 367b4db8ab1e3..ab13f52046d3e 100644 --- a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h +++ b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h @@ -506,19 +506,23 @@ static const char* const kOrtSessionOptionsCudaGqaWorkspaceMaxTotalSequenceLengt /// CUDA DynamicSparseAttention kernel switches. Each is resolved once when the kernel is created, so a change needs a new /// session, and overrides the environment variable named below. The 0/1 switches accept only "0" and "1", and the /// thresholds a non-negative decimal integer; any other value, including an empty string, fails session initialization. -/// "0" makes the INT8-cache BF16 prefill dequantize the gathered K/V tiles to BF16 instead of running FP16 tensor-core -/// MMA on them. Default "1". Overrides ORT_DSA_PREFILL_INT8_TC. +/// "1" makes the INT8-cache BF16 prefill run FP16 tensor-core MMA on the gathered K/V tiles instead of dequantizing them +/// to BF16 (opt-in; the tensor-core kernel has different numerics). Default "0". Overrides ORT_DSA_PREFILL_INT8_TC. static const char* const kOrtSessionOptionsCudaDsaPrefillInt8Tc = "ep.cuda.dsa_prefill_int8_tc"; /// "0" routes the BF16 prefill through the generic kernels. Default "1". Overrides ORT_DSA_PREFILL_KERNEL. static const char* const kOrtSessionOptionsCudaDsaPrefillKernel = "ep.cuda.dsa_prefill_kernel"; -/// "0" makes the INT8-cache BF16 decode (up to 8 query rows) use the CUDA-core grouped kernel instead of the split-K -/// tensor-core kernel. Default "1". Overrides ORT_DSA_DECODE_INT8_TC. +/// "1" makes the INT8-cache BF16 decode (up to 8 query rows) use the split-K tensor-core kernel instead of the +/// CUDA-core grouped kernel (opt-in; the tensor-core kernel accumulates in a different order and does not reproduce +/// the CUDA-core results bit for bit). Default "0". Overrides ORT_DSA_DECODE_INT8_TC. static const char* const kOrtSessionOptionsCudaDsaDecodeInt8Tc = "ep.cuda.dsa_decode_int8_tc"; -/// Number of 64-candidate tiles each CTA of the split-K tensor-core decode kernel consumes. "0" derives it from the -/// problem size (about 384 CTAs in total). Default "0". Overrides ORT_DSA_DECODE_INT8_TC_TILES_PER_SPLIT. +/// Requested number of 64-candidate tiles each CTA of the split-K tensor-core decode kernel consumes. "0" derives it from +/// the problem size and the device: the smallest power of two with which all CTAs fit one wave of the CTAs that can be +/// resident on the device. A positive value is a target: the number of splits is ceil(tiles / value), limited to 32 and +/// to the split count of the decode workspace, so a small value on a long candidate list is raised and the last split +/// may hold fewer tiles. Default "0". Overrides ORT_DSA_DECODE_INT8_TC_TILES_PER_SPLIT. static const char* const kOrtSessionOptionsCudaDsaDecodeInt8TcTilesPerSplit = "ep.cuda.dsa_decode_int8_tc_tiles_per_split"; diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc index 9bd7e62b99008..c115a8c9db26e 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc @@ -105,12 +105,12 @@ DynamicSparseAttention::DynamicSparseAttention(const OpKernelInfo& info) rotary_interleaved_ = ParseBoolAttribute(info, "rotary_interleaved", 0); use_smooth_softmax_ = ParseBoolAttribute(info, "smooth_softmax", 0); auxiliary_kv_shared_ = ParseBoolAttribute(info, "auxiliary_kv_shared", 0); - prefill_kernel_ = ReadCudaSessionOption(info, kOrtSessionOptionsCudaDsaPrefillKernel, - "ORT_DSA_PREFILL_KERNEL", 1, 0, 1) != 0; - prefill_int8_tc_ = ReadCudaSessionOption(info, kOrtSessionOptionsCudaDsaPrefillInt8Tc, - "ORT_DSA_PREFILL_INT8_TC", 1, 0, 1) != 0; - decode_int8_tc_ = ReadCudaSessionOption(info, kOrtSessionOptionsCudaDsaDecodeInt8Tc, - "ORT_DSA_DECODE_INT8_TC", 1, 0, 1) != 0; + prefill_kernel_ = ReadCudaSessionSwitch(info, kOrtSessionOptionsCudaDsaPrefillKernel, "ORT_DSA_PREFILL_KERNEL", true); + // The tensor-core INT8-cache kernels are opt-in: their numerics differ from the kernels they replace, so the default + // dispatch stays on the BF16-dequantizing prefill and the CUDA-core grouped decode kernels. + prefill_int8_tc_ = + ReadCudaSessionSwitch(info, kOrtSessionOptionsCudaDsaPrefillInt8Tc, "ORT_DSA_PREFILL_INT8_TC", false); + decode_int8_tc_ = ReadCudaSessionSwitch(info, kOrtSessionOptionsCudaDsaDecodeInt8Tc, "ORT_DSA_DECODE_INT8_TC", false); decode_int8_tc_tiles_per_split_ = ReadCudaSessionOption(info, kOrtSessionOptionsCudaDsaDecodeInt8TcTilesPerSplit, "ORT_DSA_DECODE_INT8_TC_TILES_PER_SPLIT", 0, 0, std::numeric_limits::max()); diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu index be5c6dd398fba..eaae0df5d3ab7 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu @@ -40,7 +40,6 @@ constexpr int kGroupedDecodeCandidatesPerSplit = 64; constexpr int kTargetGroupedDecodeBlocks = 128; constexpr int kGroupedQuantizedDecodeCandidatesPerSplit = 32; constexpr int kTargetGroupedQuantizedDecodeBlocks = 512; -constexpr int kTargetInt8TcDecodeCtas = 384; template __device__ __forceinline__ float ToFloat(T value) { @@ -1808,20 +1807,27 @@ int GetGroupedDecodeSplitCount(const DynamicSparseAttentionParameters& parameter return static_cast(std::min({candidate_splits, occupancy_splits, max_splits})); } -// Split count of the tensor-core INT8 decode kernel: one CTA per (row, KV head, split) and about -// kTargetInt8TcDecodeCtas CTAs in total (one wave of three CTAs per SM on 132 SMs), each consuming at least one -// 64-candidate tile. The time per CTA is dominated by a fixed latency chain, so the kernel is flat in the tiles per -// split (H200, 8 rows, 33 tiles: 1 -> 10.8 us, 2 -> 9.5, 3 -> 10.7, 4 -> 9.9, 6 -> 11.9). It never exceeds the split -// count the workspace was sized for. A positive tiles_per_split_override replaces the automatic choice. +// Split count of the tensor-core INT8 decode kernel: one CTA per (row, KV head, split), each consuming at least one +// 64-candidate tile. The automatic choice is the smallest power-of-two number of tiles per split with which all CTAs +// fit one wave of the device, resident_ctas being the SM count times the CTAs per SM that the kernel's shared-memory +// and register footprint allow (396 on an H200). A partly filled second wave costs more than longer splits do (H200 +// with the occupancy limited to one CTA per SM, 8 rows, 32 tiles: 16 splits take 18.7 us, 8 splits 14.7 us), and a +// power of two spreads the tiles evenly (H200, 16 rows, 32 tiles: 11 splits of up to 3 tiles take 20.5 us, 8 splits +// 18.5 us). The kernel is not flat in the tiles per split (H200, 8 rows, 32 tiles: 1 -> 20.4 us, 2 -> 15.0, 4 -> 14.5, +// 8 -> 19.5, 16 -> 30.9). The count never exceeds the split count the workspace was sized for. A positive +// tiles_per_split_override replaces the automatic choice. int GetInt8TcDecodeSplitCount(const DynamicSparseAttentionParameters& parameters, int workspace_splits, - int tiles_per_split_override) { + int tiles_per_split_override, int resident_ctas) { const int64_t tiles = (static_cast(parameters.max_selected) + dsa_prefill::kTile - 1) / dsa_prefill::kTile; const int64_t row_kv_heads = static_cast(parameters.batch_size) * parameters.sequence_length * parameters.kv_num_heads; - const int64_t tiles_per_split = - tiles_per_split_override > 0 - ? tiles_per_split_override - : std::max(1, (tiles * row_kv_heads + kTargetInt8TcDecodeCtas - 1) / kTargetInt8TcDecodeCtas); + int64_t tiles_per_split = tiles_per_split_override; + if (tiles_per_split <= 0) { + tiles_per_split = 1; + while (tiles_per_split * std::max(resident_ctas, 1) < tiles * row_kv_heads) { + tiles_per_split *= 2; + } + } const int64_t splits = (tiles + tiles_per_split - 1) / tiles_per_split; return static_cast(std::max(1, std::min({splits, workspace_splits, 32}))); } @@ -2148,6 +2154,7 @@ Status LaunchDynamicSparseAttention( // Decode (sequence_length <= 8) over the INT8 cache: the same tensor-core pipeline split over the candidate tiles of // each (row, KV head); the partials go through the grouped decode merge kernel. bool use_int8_tc_decode = false; + int int8_tc_decode_resident_ctas = 0; if constexpr (std::is_same_v) { if (use_grouped_decode && data.quantized_cache && data.decode_int8_tc) { int device = 0; @@ -2170,20 +2177,34 @@ Status LaunchDynamicSparseAttention( IsAligned(data.prepared_query, 16) && IsAligned(data.present_key, 16) && IsAligned(data.present_value, 16) && IsAligned(data.key_scale, 16) && IsAligned(data.value_scale, 16) && IsAligned(data.attention_workspace, 16) && - static_cast(optin_shared) >= dsa_prefill::kInt8SharedBytes; + static_cast(optin_shared) >= dsa_prefill::kInt8SplitKSharedBytes; if (use_int8_tc_decode) { // A build whose PTX is older than compute_80 holds an empty stub for the kernel. cudaFuncAttributes attributes{}; CUDA_RETURN_IF_ERROR(cudaFuncGetAttributes(&attributes, dsa_prefill::DynamicSparseAttentionInt8TcKernel)); use_int8_tc_decode = attributes.ptxVersion >= 80; } + if (use_int8_tc_decode) { + // The split count targets one wave of resident CTAs: the occupancy query needs the shared-memory opt-in. + CUDA_RETURN_IF_ERROR(cudaFuncSetAttribute( + dsa_prefill::DynamicSparseAttentionInt8TcKernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(dsa_prefill::kInt8SplitKSharedBytes))); + int multiprocessors = 0; + int blocks_per_multiprocessor = 0; + CUDA_RETURN_IF_ERROR(cudaDeviceGetAttribute(&multiprocessors, cudaDevAttrMultiProcessorCount, device)); + CUDA_RETURN_IF_ERROR(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &blocks_per_multiprocessor, dsa_prefill::DynamicSparseAttentionInt8TcKernel, dsa_prefill::kThreads, + dsa_prefill::kInt8SplitKSharedBytes)); + int8_tc_decode_resident_ctas = multiprocessors * blocks_per_multiprocessor; + } } } if (use_grouped_decode) { const int workspace_splits = GetGroupedDecodeSplitCount(parameters, data.quantized_cache); const int split_count = use_int8_tc_decode - ? GetInt8TcDecodeSplitCount(parameters, workspace_splits, data.decode_int8_tc_tiles_per_split) + ? GetInt8TcDecodeSplitCount(parameters, workspace_splits, data.decode_int8_tc_tiles_per_split, + int8_tc_decode_resident_ctas) : workspace_splits; const int64_t partial_blocks = query_blocks * split_count; const int64_t grouped_blocks = @@ -2206,13 +2227,10 @@ Status LaunchDynamicSparseAttention( GetGroupedAttentionSharedBytes(parameters, data.quantized_cache ? sizeof(float) : sizeof(T)); if (use_int8_tc_decode) { if constexpr (std::is_same_v) { - CUDA_RETURN_IF_ERROR(cudaFuncSetAttribute( - dsa_prefill::DynamicSparseAttentionInt8TcKernel, cudaFuncAttributeMaxDynamicSharedMemorySize, - static_cast(dsa_prefill::kInt8SharedBytes))); const dim3 grid(static_cast(parameters.batch_size * parameters.sequence_length), static_cast(parameters.kv_num_heads), static_cast(split_count)); dsa_prefill::DynamicSparseAttentionInt8TcKernel - <<>>( + <<>>( data.prepared_query, reinterpret_cast(data.present_key), reinterpret_cast(data.present_value), data.selected_indices, data.selected_counts, data.seqlens_k, nullptr, partial_max, partial_sum, partial_output, split_count, diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h index 24f96686efbf4..73f6e56db4231 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h @@ -51,8 +51,8 @@ struct DynamicSparseAttentionData { float* attention_workspace = nullptr; bool quantized_cache = false; bool prefill_kernel = true; - bool prefill_int8_tc = true; - bool decode_int8_tc = true; + bool prefill_int8_tc = false; + bool decode_int8_tc = false; int decode_int8_tc_tiles_per_split = 0; const float* key_scale = nullptr; const float* value_scale = nullptr; diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh index 1cda8401d6471..a0982f94a0447 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh @@ -12,13 +12,22 @@ // warp reduction and exponentials made it issue-bound, with the tensor-core pipeline below. // // The scales of the INT8 cache are per (kv head, channel) constants, so they leave the tile loop: -// * the key scale is folded into the query once per CTA: Q' = Q * key_scale (rounded once to fp16), S = Q' K_int8; -// a CTA whose Q' does not fit fp16 (or is not finite) evaluates its query heads in fp32 instead (AttendRowFp32). -// Products below 2^-25 flush to zero, which moves a score by at most 2^-10 * scale (256 channels, |K| <= 128), so -// the dispatcher selects this kernel only for scale <= kInt8TcMaxScale; -// * the value scale is applied to the output once per CTA: O = (P V_int8) * value_scale; a CTA whose value_scale -// exceeds kInt8TcMaxValueScale (so that a probability flushed below 2^-40 by the fp16 probability tile could -// still move the scaled output) is also evaluated in fp32 by AttendRowFp32. +// * the key scale is folded into the query once per CTA: Q' = Q * key_scale (rounded once to fp16), S = Q' K_int8. +// Rounding Q' moves the logit of any candidate by at most scale * 128 * sum_d |Q'_d - fp16(Q'_d)| (|K| <= 128). +// The CTA evaluates this bound for every query head and evaluates its heads in fp32 instead (AttendRowFp32) when +// it exceeds kInt8TcQueryErrorBudget, so that a query whose large terms cancel cannot amplify the rounding; the +// same happens when Q' does not fit fp16 or is not finite. The bound includes the flush of products below 2^-25, +// and the dispatcher selects this kernel only for scale <= kInt8TcMaxScale, which keeps the bound far below the +// budget for ordinary data; +// * the value scale is applied to the output once per CTA: O = (P V_int8) * value_scale. A probability below 2^-29 +// (more than kInt8TcTailDepth nats below the running maximum) leaves the normal range of the fp16 probability tile +// and loses precision (below 2^-40 it flushes to zero), which a large value scale could amplify. A CTA with a +// channel |value_scale| > kInt8TcMaxValueScale therefore tracks how far below the running maximum its valid +// candidates are and, only if one of them is deeper than kInt8TcTailDepth, discards the tile loop and evaluates +// its heads in fp32 as well. A value scale at or below kInt8TcMaxValueScale never takes this path. Above it, a CTA +// takes it when one of its heads has a candidate kInt8TcTailDepth nats below its running maximum, about 7 standard +// deviations of Gaussian logits over 2048 candidates: a fraction of the CTAs when the standard deviation is about +// 3 nats, practically none for 1 nat. The documentation lists measurements. // The gathered int8 rows stay int8 in shared memory (cp.async, double buffered, one tile ahead). A fragment read with // ldmatrix is converted to fp16 exactly in registers (1024 + (byte ^ 0x80) - 1152, "magic number" conversion) and // consumed by fp16 MMAs with fp32 accumulation, so the per-tile dequantization pass, its shared memory round trip and @@ -26,6 +35,13 @@ // per tile are the row maxima and the fp16 probability tile, which holds kInt8ProbScale * p so that the tail of the // softmax stays out of the fp16 subnormal range. // +// Scores and the running maximum are kept in the natural logarithm domain and only the difference to the new maximum +// is scaled by log2(e) for ex2: the difference of two nearby large logits is then exact, which the product with log2(e) +// would round away. Decode (kSplitK) splits the probability tile into an fp16 high part and an fp16 residual (about +// 22 significant bits together), because the signed sum of large value codes can cancel to a value far below the 2^-11 +// rounding of a single fp16 probability; both parts multiply the same converted V fragments. The prefill variant keeps +// the single fp16 tile (as does the BF16-dequantizing kernel, whose bfloat16 tile is coarser still). +// // The channel (reduction) dimension of S = Q' K^T is permuted so that one 4-byte ldmatrix element of a K row supplies // both B registers of an MMA k step: slot 2t+e <-> channel 16j + 4t + e and slot 2t+8+e <-> channel 16j + 4t + 2 + e, // with the same permutation applied to the A fragments of Q'. Likewise ldmatrix.trans on the V tile delivers channels @@ -52,22 +68,35 @@ constexpr int kIdxStages = 3; // The fp16 probability tile holds kInt8ProbScale * p for p in [0, 1]. Without the factor a candidate whose probability // is below 2^-25 rounds to zero although a large value scale can make its contribution matter; with it the flush point // drops to 2^-40 and the tile stays below the fp16 maximum. The PV accumulator carries the factor until the epilogue -// removes it; the softmax denominator is accumulated in fp32 without it. When any channel has -// |value_scale| > kInt8TcMaxValueScale (so that |p * V * value_scale| for p <= 2^-40 can exceed 2^-17), the CTA takes -// AttendRowFp32 instead. +// removes it; the softmax denominator is accumulated in fp32 without it. A probability below 2^-29 is a subnormal +// fp16 value of the tile (kInt8TcTailDepth nats below the running maximum). When some channel has +// |value_scale| > kInt8TcMaxValueScale (so that |p * V * value_scale| for p <= 2^-40 can exceed 2^-17) and a valid +// candidate is that deep, the CTA takes AttendRowFp32 instead. constexpr float kInt8ProbScale = 32768.0f; // 2^15 constexpr float kInt8ProbInverseScale = 1.0f / kInt8ProbScale; constexpr float kInt8TcMaxValueScale = 65536.0f; // 2^16 +constexpr float kInt8TcTailDepth = 20.0f; // nats; 2^-29 is 20.1 -// Largest softmax scale served by the kernel below (the op validates it as finite and non-negative). Rounding -// Q * key_scale to fp16 flushes products below 2^-25 to zero, which moves a dot product by at most -// 256 * 128 * 2^-25 = 2^-10 before the scale multiplies it; a larger scale would amplify that error, so the dispatcher -// uses the BF16-dequantizing prefill / CUDA-core decode kernels for it. +// Largest softmax scale served by the kernel below (the op validates it as finite and non-negative). The error +// bound on Q * key_scale rounded to fp16 grows linearly with it, so a larger scale would send ordinary data to the +// fp32 fallback; the dispatcher uses the BF16-dequantizing prefill / CUDA-core decode kernels for it instead. constexpr float kInt8TcMaxScale = 1.0f; +// Largest bound on the change of a logit caused by rounding Q * key_scale to fp16 that the tensor-core path accepts: +// scale * 128 * sum_d |Q'_d - fp16(Q'_d)| of the query head (see the top). For a normal-range Q' the bound is at most +// 2^-4 * scale * sum_d |Q'_d|, an eighth of what rounding the cache to INT8 can move the same logit, so ordinary +// queries stay far below the budget, while a query whose large terms cancel to a small logit (the INT8 error of an +// exactly representable cache is zero there) exceeds it by orders of magnitude and is evaluated in fp32. The bound is +// about 0.015 for unit-RMS queries scaled by the q-norm weights and the key scales of the Qwen3.8-Flash-Next DSA +// layers (docs/contrib_ops/cuda/dynamic_sparse_attention.md lists how it grows with the query magnitude). +constexpr float kInt8TcQueryErrorBudget = 0.125f; + // sK8[2] + sV8[2] (sQ aliases sK8[0] before the tile loop) + candidate indices + row maxima / sums + probabilities. +constexpr size_t kInt8PBytes = kMaxGroup * kInt8PStride * sizeof(__half); constexpr size_t kInt8SharedBytes = 4 * static_cast(kInt8TileBytes) + kIdxStages * kTile * sizeof(int32_t) + - kMaxGroup * 4 * sizeof(float) + kMaxGroup * kInt8PStride * sizeof(__half); + kMaxGroup * 4 * sizeof(float) + kInt8PBytes; +// The split-K (decode) variant adds the tile of fp16 probability residuals. +constexpr size_t kInt8SplitKSharedBytes = kInt8SharedBytes + kInt8PBytes; #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 __device__ __forceinline__ void CpAsync16Zfill(uint32_t smem_dst, const void* gmem_src, bool valid) { @@ -129,15 +158,17 @@ __device__ __forceinline__ void Int8x4ToHalf2Strided(uint32_t w, uint32_t& even, odd = Half2MinusBias(b); } -// FP32 evaluation of the query heads of one (row, KV head), one warp per head and eight channels per lane, for rows -// whose Q * key_scale leaves the FP16 range of the tensor-core path or whose value_scale exceeds kInt8TcMaxValueScale. -// Clipping Q * key_scale would break the cancellation between large terms, and a larger value_scale can amplify -// probabilities below the 2^-40 FP16 probability-tile flush boundary. `query` and `output` point at the first head of +// FP32 evaluation of the query heads of one (row, KV head), one warp per head and eight channels per lane, for CTAs +// the tensor-core path cannot serve accurately: Q * key_scale leaves the FP16 range or its rounding error bound exceeds +// kInt8TcQueryErrorBudget (clipping or rounding would break the cancellation between large terms), or a value_scale +// above kInt8TcMaxValueScale meets a probability below the normal range of the FP16 probability tile (the tile flushes +// below 2^-40 and a large value scale can amplify that). `query` and `output` point at the first head of // the group; `key_scale` and `value_scale` at the KV head. The candidates [begin, end) of `row_indices` are consumed: // all of them for prefill, the range of the CTA for a decode split. With kSplitK the un-normalized partial (running // maximum in the natural logarithm domain, running sum, value-scaled accumulator) of block // partial_block + head * split_count goes to partial_*, as the main path does; otherwise the normalized result goes to -// output. +// output. Every candidate is read from global memory once per head and handled serially, so this path is one to two +// orders of magnitude slower than the tensor-core one. template __device__ __noinline__ void AttendRowFp32(const __nv_bfloat16* __restrict__ query, const int8_t* __restrict__ key, const int8_t* __restrict__ value, const float* __restrict__ key_scale, @@ -219,7 +250,7 @@ __device__ __noinline__ void AttendRowFp32(const __nv_bfloat16* __restrict__ que #endif // grid = (batch * sequence_length, kv_num_heads) or, with kSplitK, (batch * sequence_length, kv_num_heads, -// split_count); block = kThreads, dynamic smem = kInt8SharedBytes. +// split_count); block = kThreads, dynamic smem = kInt8SharedBytes (kInt8SplitKSharedBytes with kSplitK). // query: [rows, num_heads, 256] (already normalized / rotated); key, value: int8 [batch, kv_heads, capacity, 256]; // key_scale, value_scale: float [kv_heads, 256]. The normalized result goes to output; with kSplitK the partials // go to partial_max / partial_sum [rows * num_heads * split_count] and partial_output [.., 256] instead, indexed @@ -253,6 +284,8 @@ __global__ void __launch_bounds__(kThreads, 3) DynamicSparseAttentionInt8TcKerne float* sM = reinterpret_cast(sIdx + kIdxStages * kTile); __half* sP = reinterpret_cast<__half*>(sM + kMaxGroup * 4); __half* sQ = reinterpret_cast<__half*>(smem_raw); + // kSplitK only: fp16 residual of the probability tile, right behind sP (see kInt8SplitKSharedBytes). + [[maybe_unused]] __half* sPResidual = sP + kMaxGroup * kInt8PStride; const int tid = threadIdx.x; const int lane = tid & 31; @@ -295,30 +328,41 @@ __global__ void __launch_bounds__(kThreads, 3) DynamicSparseAttentionInt8TcKerne } } - // Q' = Q * key_scale as fp16; the padding rows are zero so their (ignored) scores stay finite. A row whose Q' does - // not fit fp16 (including inf / NaN) or whose value_scale exceeds kInt8TcMaxValueScale is evaluated in fp32 by the - // whole CTA instead. Underflow to fp16 subnormals changes a product by at most 2^-25, i.e. a score by at most - // 2^-10 * scale, which the dispatcher keeps at the level of the fp16 rounding by requiring scale <= kInt8TcMaxScale. + // Q' = Q * key_scale as fp16; the padding rows are zero so their (ignored) scores stay finite. The whole CTA is + // evaluated in fp32 by AttendRowFp32 instead when Q' of a head does not fit fp16 (including inf / NaN), when the bound + // scale * 128 * sum_d |Q'_d - fp16(Q'_d)| on the rounding error of a logit exceeds kInt8TcQueryErrorBudget for a + // head, or when a value_scale is not finite. A finite value_scale above kInt8TcMaxValueScale only arms the check of + // the probability depth after the tile loop. const __nv_bfloat16* query_group = query + (static_cast(row) * num_heads + kv_head * group) * kHeadDim; bool needs_fp32 = false; + bool large_value_scale = false; { const int chunk = tid & 31; const float4 k0 = *reinterpret_cast(key_scale + kv_head * kHeadDim + chunk * 8); const float4 k1 = *reinterpret_cast(key_scale + kv_head * kHeadDim + chunk * 8 + 4); if (tid < kHeadDim / 4) { const float4 v = *reinterpret_cast(value_scale + kv_head * kHeadDim + tid * 4); - needs_fp32 = !(fabsf(v.x) <= kInt8TcMaxValueScale) || !(fabsf(v.y) <= kInt8TcMaxValueScale) || - !(fabsf(v.z) <= kInt8TcMaxValueScale) || !(fabsf(v.w) <= kInt8TcMaxValueScale); + needs_fp32 = !(fabsf(v.x) <= FLT_MAX) || !(fabsf(v.y) <= FLT_MAX) || !(fabsf(v.z) <= FLT_MAX) || + !(fabsf(v.w) <= FLT_MAX); + large_value_scale = fabsf(v.x) > kInt8TcMaxValueScale || fabsf(v.y) > kInt8TcMaxValueScale || + fabsf(v.z) > kInt8TcMaxValueScale || fabsf(v.w) > kInt8TcMaxValueScale; } - auto to_half2 = [&needs_fp32](float x, float y) { + // Sum of |Q' - fp16(Q')| over the eight channels of this lane (NaN when Q' is not finite or leaves fp16). + float rounding_error = 0.0f; + auto to_half2 = [&needs_fp32, &rounding_error](float x, float y) { needs_fp32 |= !(fabsf(x) <= 65504.0f) || !(fabsf(y) <= 65504.0f); - return __floats2half2_rn(x, y); + const __half2 h = __floats2half2_rn(x, y); + const float2 rounded = __half22float2(h); + rounding_error += fabsf(x - rounded.x) + fabsf(y - rounded.y); + return h; }; + const float error_scale = scale * 128.0f; #pragma unroll for (int i = 0; i < kMaxGroup / 4; ++i) { const int r = (tid >> 5) + 4 * i; uint4 packed = make_uint4(0, 0, 0, 0); if (r < group) { + rounding_error = 0.0f; const uint4 q = *reinterpret_cast(query_group + static_cast(r) * kHeadDim + chunk * 8); const float2 q01 = __bfloat1622float2(*reinterpret_cast(&q.x)); const float2 q23 = __bfloat1622float2(*reinterpret_cast(&q.y)); @@ -332,12 +376,19 @@ __global__ void __launch_bounds__(kThreads, 3) DynamicSparseAttentionInt8TcKerne packed.y = *reinterpret_cast(&h23); packed.z = *reinterpret_cast(&h45); packed.w = *reinterpret_cast(&h67); + // The warp covers all 256 channels of row r (r < group is uniform across it): a butterfly sum gives the + // bound on the logit error of head r in every lane. +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + rounding_error += __shfl_xor_sync(0xffffffffu, rounding_error, offset); + } + needs_fp32 |= !(rounding_error * error_scale <= kInt8TcQueryErrorBudget); } *reinterpret_cast(sQ + r * kInt8QStride + chunk * 8) = packed; } } - if (__syncthreads_or(needs_fp32)) { - // The condition is uniform across the CTA and no cp.async is in flight yet. + // The condition is uniform across the CTA and no cp.async is in flight yet. + auto attend_fp32 = [&]() { const float* kv_key_scale = key_scale + kv_head * kHeadDim; const float* kv_value_scale = value_scale + kv_head * kHeadDim; if constexpr (kSplitK) { @@ -352,6 +403,9 @@ __global__ void __launch_bounds__(kThreads, 3) DynamicSparseAttentionInt8TcKerne output + (static_cast(row) * num_heads + kv_head * group) * kHeadDim, nullptr, nullptr, nullptr, 0, 1, warp, lane); } + }; + if (__syncthreads_or(needs_fp32)) { + attend_fp32(); return; } @@ -367,7 +421,8 @@ __global__ void __launch_bounds__(kThreads, 3) DynamicSparseAttentionInt8TcKerne qf[j][2] = lo.y; qf[j][3] = hi.y; } - __syncthreads(); // sQ aliases sK8[0] + // sQ aliases sK8[0]: this barrier also tells the CTA whether the probability depth has to be checked (see below). + const bool check_tail_depth = __syncthreads_or(large_value_scale); auto resolve_index = [&](int raw) -> int { return (raw >= 0 && raw < total && raw <= query_position && raw < capacity) ? raw : -1; @@ -384,6 +439,7 @@ __global__ void __launch_bounds__(kThreads, 3) DynamicSparseAttentionInt8TcKerne const uint32_t k_smem_base = static_cast(__cvta_generic_to_shared(sK8)); const uint32_t v_smem_base = static_cast(__cvta_generic_to_shared(sV8)); const uint32_t p_smem_base = static_cast(__cvta_generic_to_shared(sP)); + [[maybe_unused]] const uint32_t p_residual_smem_base = p_smem_base + static_cast(kInt8PBytes); const int load_chunk = tid & 15; // 16 B per lane: 16 lanes per candidate row, 8 rows per pass, 8 passes per tile (K and V). @@ -410,13 +466,16 @@ __global__ void __launch_bounds__(kThreads, 3) DynamicSparseAttentionInt8TcKerne for (int n = 0; n < 8; ++n) { acc_o[n][0] = acc_o[n][1] = acc_o[n][2] = acc_o[n][3] = 0.0f; } - // Softmax state of rows r0 / r0 + 8: the running maximum is identical in all threads of the row, the running sum only - // covers this thread's columns (merged across threads and warps after the loop). + // Softmax state of rows r0 / r0 + 8: the running maximum (natural logarithm domain) is identical in all threads of + // the row, the running sum only covers this thread's columns (merged across threads and warps after the loop). float run_max0 = -FLT_MAX; float run_max1 = -FLT_MAX; float run_sum0 = 0.0f; float run_sum1 = 0.0f; - const float scale_log2 = scale * 1.44269504088896341f; + // Most negative (score - running maximum) of a valid candidate seen by this thread; only tracked when a value scale + // is large. + float lowest_depth = 0.0f; + constexpr float kLog2e = 1.44269504088896341f; // Shared memory byte addresses that do not change from tile to tile. const uint32_t k_lane_offset = @@ -470,7 +529,7 @@ __global__ void __launch_bounds__(kThreads, 3) DynamicSparseAttentionInt8TcKerne } } - // Scores in the log2 domain; invalid candidates are masked. + // Scores in the natural logarithm domain; invalid candidates are masked. const int cbase = warp * 16 + 2 * t; const int2 ia = *reinterpret_cast(sIdx + idx_cur * kTile + cbase); const int2 ib = *reinterpret_cast(sIdx + idx_cur * kTile + cbase + 8); @@ -478,14 +537,14 @@ __global__ void __launch_bounds__(kThreads, 3) DynamicSparseAttentionInt8TcKerne const bool va1 = ia.y >= 0; const bool vb0 = ib.x >= 0; const bool vb1 = ib.y >= 0; - const float s00 = va0 ? acc_s[0][0] * scale_log2 : -FLT_MAX; - const float s01 = va1 ? acc_s[0][1] * scale_log2 : -FLT_MAX; - const float s02 = vb0 ? acc_s[1][0] * scale_log2 : -FLT_MAX; - const float s03 = vb1 ? acc_s[1][1] * scale_log2 : -FLT_MAX; - const float s10 = va0 ? acc_s[0][2] * scale_log2 : -FLT_MAX; - const float s11 = va1 ? acc_s[0][3] * scale_log2 : -FLT_MAX; - const float s12 = vb0 ? acc_s[1][2] * scale_log2 : -FLT_MAX; - const float s13 = vb1 ? acc_s[1][3] * scale_log2 : -FLT_MAX; + const float s00 = va0 ? acc_s[0][0] * scale : -FLT_MAX; + const float s01 = va1 ? acc_s[0][1] * scale : -FLT_MAX; + const float s02 = vb0 ? acc_s[1][0] * scale : -FLT_MAX; + const float s03 = vb1 ? acc_s[1][1] * scale : -FLT_MAX; + const float s10 = va0 ? acc_s[0][2] * scale : -FLT_MAX; + const float s11 = va1 ? acc_s[0][3] * scale : -FLT_MAX; + const float s12 = vb0 ? acc_s[1][2] * scale : -FLT_MAX; + const float s13 = vb1 ? acc_s[1][3] * scale : -FLT_MAX; float m0 = fmaxf(fmaxf(s00, s01), fmaxf(s02, s03)); float m1 = fmaxf(fmaxf(s10, s11), fmaxf(s12, s13)); m0 = fmaxf(m0, __shfl_xor_sync(0xffffffffu, m0, 1)); @@ -502,31 +561,69 @@ __global__ void __launch_bounds__(kThreads, 3) DynamicSparseAttentionInt8TcKerne const float4 mx1 = *reinterpret_cast(sM + (r0 + 8) * 4); const float new_max0 = fmaxf(run_max0, fmaxf(fmaxf(mx0.x, mx0.y), fmaxf(mx0.z, mx0.w))); const float new_max1 = fmaxf(run_max1, fmaxf(fmaxf(mx1.x, mx1.y), fmaxf(mx1.z, mx1.w))); - const float alpha0 = FastExp2(run_max0 - new_max0); - const float alpha1 = FastExp2(run_max1 - new_max1); - const float p00 = va0 ? FastExp2(s00 - new_max0) : 0.0f; - const float p01 = va1 ? FastExp2(s01 - new_max0) : 0.0f; - const float p02 = vb0 ? FastExp2(s02 - new_max0) : 0.0f; - const float p03 = vb1 ? FastExp2(s03 - new_max0) : 0.0f; - const float p10 = va0 ? FastExp2(s10 - new_max1) : 0.0f; - const float p11 = va1 ? FastExp2(s11 - new_max1) : 0.0f; - const float p12 = vb0 ? FastExp2(s12 - new_max1) : 0.0f; - const float p13 = vb1 ? FastExp2(s13 - new_max1) : 0.0f; + // The differences to the new maximum are formed in the natural logarithm domain, where nearby large logits are + // still distinct, and only then scaled by log2(e) for ex2. + const float alpha0 = FastExp2((run_max0 - new_max0) * kLog2e); + const float alpha1 = FastExp2((run_max1 - new_max1) * kLog2e); + const float d00 = s00 - new_max0; + const float d01 = s01 - new_max0; + const float d02 = s02 - new_max0; + const float d03 = s03 - new_max0; + const float d10 = s10 - new_max1; + const float d11 = s11 - new_max1; + const float d12 = s12 - new_max1; + const float d13 = s13 - new_max1; + const float p00 = va0 ? FastExp2(d00 * kLog2e) : 0.0f; + const float p01 = va1 ? FastExp2(d01 * kLog2e) : 0.0f; + const float p02 = vb0 ? FastExp2(d02 * kLog2e) : 0.0f; + const float p03 = vb1 ? FastExp2(d03 * kLog2e) : 0.0f; + const float p10 = va0 ? FastExp2(d10 * kLog2e) : 0.0f; + const float p11 = va1 ? FastExp2(d11 * kLog2e) : 0.0f; + const float p12 = vb0 ? FastExp2(d12 * kLog2e) : 0.0f; + const float p13 = vb1 ? FastExp2(d13 * kLog2e) : 0.0f; run_sum0 = run_sum0 * alpha0 + ((p00 + p01) + (p02 + p03)); run_sum1 = run_sum1 * alpha1 + ((p10 + p11) + (p12 + p13)); run_max0 = new_max0; run_max1 = new_max1; + if (check_tail_depth) { + const float low0 = fminf(fminf(va0 ? d00 : 0.0f, va1 ? d01 : 0.0f), fminf(vb0 ? d02 : 0.0f, vb1 ? d03 : 0.0f)); + const float low1 = fminf(fminf(va0 ? d10 : 0.0f, va1 ? d11 : 0.0f), fminf(vb0 ? d12 : 0.0f, vb1 ? d13 : 0.0f)); + lowest_depth = fminf(lowest_depth, fminf(low0, low1)); + } { - const __half2 h00 = __floats2half2_rn(p00 * kInt8ProbScale, p01 * kInt8ProbScale); - const __half2 h01 = __floats2half2_rn(p02 * kInt8ProbScale, p03 * kInt8ProbScale); - const __half2 h10 = __floats2half2_rn(p10 * kInt8ProbScale, p11 * kInt8ProbScale); - const __half2 h11 = __floats2half2_rn(p12 * kInt8ProbScale, p13 * kInt8ProbScale); + const float x00 = p00 * kInt8ProbScale; + const float x01 = p01 * kInt8ProbScale; + const float x02 = p02 * kInt8ProbScale; + const float x03 = p03 * kInt8ProbScale; + const float x10 = p10 * kInt8ProbScale; + const float x11 = p11 * kInt8ProbScale; + const float x12 = p12 * kInt8ProbScale; + const float x13 = p13 * kInt8ProbScale; + const __half2 h00 = __floats2half2_rn(x00, x01); + const __half2 h01 = __floats2half2_rn(x02, x03); + const __half2 h10 = __floats2half2_rn(x10, x11); + const __half2 h11 = __floats2half2_rn(x12, x13); *reinterpret_cast<__half2*>(sP + r0 * kInt8PStride + cbase) = h00; *reinterpret_cast<__half2*>(sP + r0 * kInt8PStride + cbase + 8) = h01; *reinterpret_cast<__half2*>(sP + (r0 + 8) * kInt8PStride + cbase) = h10; *reinterpret_cast<__half2*>(sP + (r0 + 8) * kInt8PStride + cbase + 8) = h11; + if constexpr (kSplitK) { + // What the fp16 tile lost: P = P_hi + P_residual keeps about 22 significant bits. + const float2 f00 = __half22float2(h00); + const float2 f01 = __half22float2(h01); + const float2 f10 = __half22float2(h10); + const float2 f11 = __half22float2(h11); + *reinterpret_cast<__half2*>(sPResidual + r0 * kInt8PStride + cbase) = + __floats2half2_rn(x00 - f00.x, x01 - f00.y); + *reinterpret_cast<__half2*>(sPResidual + r0 * kInt8PStride + cbase + 8) = + __floats2half2_rn(x02 - f01.x, x03 - f01.y); + *reinterpret_cast<__half2*>(sPResidual + (r0 + 8) * kInt8PStride + cbase) = + __floats2half2_rn(x10 - f10.x, x11 - f10.y); + *reinterpret_cast<__half2*>(sPResidual + (r0 + 8) * kInt8PStride + cbase + 8) = + __floats2half2_rn(x12 - f11.x, x13 - f11.y); + } } - __syncthreads(); // (3) the probability tile is visible + __syncthreads(); // (3) the probability tile (and its residual) is visible // The maximum rarely changes after the first tiles: skip the accumulator rescale when the whole warp is unchanged. if (__any_sync(0xffffffffu, alpha0 != 1.0f || alpha1 != 1.0f)) { @@ -539,13 +636,18 @@ __global__ void __launch_bounds__(kThreads, 3) DynamicSparseAttentionInt8TcKerne } } - // O += P V for this warp's 64 output channels (four 16-channel chunks, each an even and an odd n tile). + // O += P V for this warp's 64 output channels (four 16-channel chunks, each an even and an odd n tile). With kSplitK + // the residual tile multiplies the same converted V fragments. { const uint32_t v_addr = v_smem_base + stage * kInt8TileBytes + v_lane_offset; #pragma unroll for (int kk = 0; kk < kTile / 16; ++kk) { uint32_t a[4]; LdmatrixX4Addr(a, p_smem_base + p_lane_offset + kk * 32); + uint32_t a_residual[4] = {0, 0, 0, 0}; + if constexpr (kSplitK) { + LdmatrixX4Addr(a_residual, p_residual_smem_base + p_lane_offset + kk * 32); + } #pragma unroll for (int cp = 0; cp < 2; ++cp) { uint32_t r[4]; @@ -559,10 +661,18 @@ __global__ void __launch_bounds__(kThreads, 3) DynamicSparseAttentionInt8TcKerne Int8x4ToHalf2Strided(r[1], e1, o1); MmaF16(acc_o[4 * cp], a, e0, e1); MmaF16(acc_o[4 * cp + 1], a, o0, o1); + if constexpr (kSplitK) { + MmaF16(acc_o[4 * cp], a_residual, e0, e1); + MmaF16(acc_o[4 * cp + 1], a_residual, o0, o1); + } Int8x4ToHalf2Strided(r[2], e0, o0); Int8x4ToHalf2Strided(r[3], e1, o1); MmaF16(acc_o[4 * cp + 2], a, e0, e1); MmaF16(acc_o[4 * cp + 3], a, o0, o1); + if constexpr (kSplitK) { + MmaF16(acc_o[4 * cp + 2], a_residual, e0, e1); + MmaF16(acc_o[4 * cp + 3], a_residual, o0, o1); + } } } } @@ -575,6 +685,14 @@ __global__ void __launch_bounds__(kThreads, 3) DynamicSparseAttentionInt8TcKerne idx_next2 = idx_next2 + 1 == kIdxStages ? 0 : idx_next2 + 1; } + // A large value scale amplifies probabilities that the fp16 tile represents with a large relative error or not at + // all (below 2^-29). The condition is uniform across the CTA and no cp.async copy is in flight here; the tile loop + // wrote nothing to global memory, so the fp32 path simply starts over from the inputs. + if (check_tail_depth && __syncthreads_or(lowest_depth < -kInt8TcTailDepth)) { + attend_fp32(); + return; + } + // Merge the per-thread running sums: across the four threads of a row, then across the four warps. float l0 = run_sum0; float l1 = run_sum1; @@ -599,36 +717,38 @@ __global__ void __launch_bounds__(kThreads, 3) DynamicSparseAttentionInt8TcKerne if constexpr (kSplitK) { // Un-normalized partial of this split in the natural logarithm domain, as the merge kernel expects. Thread t owns // channels 64 * warp + 16 * chunk + 4t + {0, 1, 2, 3} (see below), i.e. one float4 per chunk and row. - constexpr float kLn2 = 0.693147180559945309f; const int split = static_cast(blockIdx.z); const int partial_block0 = (row * num_heads + kv_head * group + r0) * split_count + split; const int partial_block1 = partial_block0 + 8 * split_count; if (warp == 0 && t == 0) { if (r0 < group) { - partial_max[partial_block0] = run_max0 * kLn2; + partial_max[partial_block0] = run_max0; partial_sum[partial_block0] = sum0; } if (r0 + 8 < group) { - partial_max[partial_block1] = run_max1 * kLn2; + partial_max[partial_block1] = run_max1; partial_sum[partial_block1] = sum1; } } #pragma unroll for (int chunk = 0; chunk < 4; ++chunk) { const int dim = warp * 64 + chunk * 16 + t * 4; - // acc_o carries kInt8ProbScale (the partial_sum above does not), which the value scale takes out. - const float4 vs_raw = *reinterpret_cast(value_scale + kv_head * kHeadDim + dim); - const float4 vs = make_float4(vs_raw.x * kInt8ProbInverseScale, vs_raw.y * kInt8ProbInverseScale, - vs_raw.z * kInt8ProbInverseScale, vs_raw.w * kInt8ProbInverseScale); + // acc_o carries kInt8ProbScale (the partial_sum above does not). Removing it before the multiplication by the + // value scale keeps a tiny scale (or the product) from underflowing in the factor. + const float4 vs = *reinterpret_cast(value_scale + kv_head * kHeadDim + dim); if (r0 < group) { *reinterpret_cast(partial_output + static_cast(partial_block0) * kHeadDim + dim) = - make_float4(acc_o[2 * chunk][0] * vs.x, acc_o[2 * chunk + 1][0] * vs.y, acc_o[2 * chunk][1] * vs.z, - acc_o[2 * chunk + 1][1] * vs.w); + make_float4((acc_o[2 * chunk][0] * kInt8ProbInverseScale) * vs.x, + (acc_o[2 * chunk + 1][0] * kInt8ProbInverseScale) * vs.y, + (acc_o[2 * chunk][1] * kInt8ProbInverseScale) * vs.z, + (acc_o[2 * chunk + 1][1] * kInt8ProbInverseScale) * vs.w); } if (r0 + 8 < group) { *reinterpret_cast(partial_output + static_cast(partial_block1) * kHeadDim + dim) = - make_float4(acc_o[2 * chunk][2] * vs.x, acc_o[2 * chunk + 1][2] * vs.y, acc_o[2 * chunk][3] * vs.z, - acc_o[2 * chunk + 1][3] * vs.w); + make_float4((acc_o[2 * chunk][2] * kInt8ProbInverseScale) * vs.x, + (acc_o[2 * chunk + 1][2] * kInt8ProbInverseScale) * vs.y, + (acc_o[2 * chunk][3] * kInt8ProbInverseScale) * vs.z, + (acc_o[2 * chunk + 1][3] * kInt8ProbInverseScale) * vs.w); } } } else { diff --git a/onnxruntime/contrib_ops/cuda/cuda_session_options.h b/onnxruntime/contrib_ops/cuda/cuda_session_options.h index 93f6ffcb64876..0f2f06bbe0e6d 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_session_options.h +++ b/onnxruntime/contrib_ops/cuda/cuda_session_options.h @@ -13,13 +13,16 @@ namespace onnxruntime::contrib::cuda { -// An explicit configured value takes precedence over the environment. Throws unless the result is an integer in -// [min_value, max_value]. +// An explicit configured value takes precedence over the environment. +inline std::string ResolveCudaSessionOptionText(const std::optional& configured, const char* env, + const std::string& default_text) { + return configured.has_value() ? *configured : ParseEnvironmentVariableWithDefault(env, default_text); +} + +// Integer threshold: throws unless the result is an integer in [min_value, max_value]. inline int ResolveCudaSessionOption(const std::optional& configured, const char* key, const char* env, int default_value, int min_value, int max_value) { - const std::string value = configured.has_value() - ? *configured - : ParseEnvironmentVariableWithDefault(env, std::to_string(default_value)); + const std::string value = ResolveCudaSessionOptionText(configured, env, std::to_string(default_value)); int parsed = 0; const auto result = std::from_chars(value.data(), value.data() + value.size(), parsed); ORT_ENFORCE(result.ec == std::errc{} && result.ptr == value.data() + value.size() && @@ -28,6 +31,14 @@ inline int ResolveCudaSessionOption(const std::optional& configured return parsed; } +// 0/1 switch: throws unless the result is exactly "0" or "1" (an integer parse would also take "00", "01" or "-0"). +inline bool ResolveCudaSessionSwitch(const std::optional& configured, const char* key, const char* env, + bool default_value) { + const std::string value = ResolveCudaSessionOptionText(configured, env, default_value ? "1" : "0"); + ORT_ENFORCE(value == "0" || value == "1", key, " (", env, ") must be exactly \"0\" or \"1\", got '", value, "'."); + return value == "1"; +} + // Call only from kernel construction. Explicit session configuration takes precedence over the environment. inline int ReadCudaSessionOption(const OpKernelInfo& info, const char* key, const char* env, int default_value, int min_value, int max_value) { @@ -35,4 +46,8 @@ inline int ReadCudaSessionOption(const OpKernelInfo& info, const char* key, cons max_value); } +inline bool ReadCudaSessionSwitch(const OpKernelInfo& info, const char* key, const char* env, bool default_value) { + return ResolveCudaSessionSwitch(info.GetConfigOptions().GetConfigEntry(key), key, env, default_value); +} + } // namespace onnxruntime::contrib::cuda From bd24bb8503f59562e8d3a70932cc0637df7a97ed Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 31 Oct 2026 20:42:58 +0000 Subject: [PATCH 11/14] [CUDA] DynamicSparseAttention tests: FP64 reference, rounding and cancellation cases, strict switch parsing Reference and tolerances: - ComputeSelectedOnlyPrefillExpectedOutput rounds the inputs the way the operator receives them (BF16 query and floating-point cache, exact INT8 codes with float scales) and evaluates the scores, the max-subtracted weights and the sums in double precision, so that a large common logit or a tiny weight neither overflows nor underflows the reference. - DynamicSparseAttentionCase carries its own output tolerance (absolute and relative) instead of a fixed 0.005. - The two-candidate INT8 fixture takes a second candidate position, a list length that puts the candidates in the first and the last of three 64-candidate tiles (two active decode splits), several tiles-per-split values, a one-bfloat16-ulp check and a hook that customizes the query, keys and scales. New deterministic cases (the output follows from a few numbers; prefill and decode, with the candidates in one tile and in the first and last of three): - Int8TensorCoreQueryFoldingRoundingError_CUDA: Q * key_scale is exact in fp32 but not an fp16 number; rounding it cancels every channel pair of a logit gap of 16 (output 0.5 instead of 1 - 1.1e-7). - Int8TensorCoreDecodeSignedValueCancellation_CUDA: two probabilities 0.5 +- 2^-14 that one fp16 tile cannot tell apart, with signed value codes that cancel. Decode only: prefill keeps a single 16-bit probability tile. - Int8TensorCoreLargeCommonLogit_CUDA: two logits near 7.5e5 whose gap a base-2 conversion loses to the float spacing; the natural-log domain keeps it. - Int8TensorCoreTinyValueScale_CUDA: value_scale 2^-135 gives a subnormal that must reach the BF16 output, in a list that leaves two splits empty. - Int8TensorCoreDecodeTilesPerSplitCaps_CUDA: 32, 33 and 128 tiles with 0, 1, 2, 3 and 64 tiles per split around the cap of 32 splits and the workspace; every candidate must reach the merge exactly once. New random cases against the FP64 reference: - Int8TensorCoreRandomQueryFolding_CUDA: key scales and queries whose folded products are not fp16 numbers, at query multipliers 1 (tensor-core scores) and 16 (error budget exceeded, FP32 scores). - Int8TensorCoreLargeValueScaleRandom_CUDA: value scales between 2^14 and 2^17 with score spreads of about 0.7, 11 and 22 nats, so that the tail check sees both outcomes on random data. Session options: PrefillSessionOptions_CUDA and DecodeSessionOptions_CUDA reject "2", "-1", "", "1x", "true", "00", "01", "-0", "000001", "+1", " 1" and "1 " as session options of the 0/1 switches (ep.cuda.dsa_prefill_kernel, ep.cuda.dsa_prefill_int8_tc, ep.cuda.dsa_decode_int8_tc), and "01", "-0" and "+1" in ORT_DSA_PREFILL_INT8_TC and ORT_DSA_DECODE_INT8_TC. Against the provider built from the previous commit (ceed60271c), these fail: QueryFoldingRoundingError, DecodeSignedValueCancellation, LargeCommonLogit, TinyValueScale, RandomQueryFolding, PrefillSessionOptions and DecodeSessionOptions. DecodeTilesPerSplitCaps and LargeValueScaleRandom pass on both; they pin the split cap and the tail check on random data. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../dynamic_sparse_attention_op_test.cc | 385 ++++++++++++++++-- 1 file changed, 346 insertions(+), 39 deletions(-) diff --git a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc index a80c9440cc2ff..a737731f104e3 100644 --- a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc @@ -4,6 +4,7 @@ #include #include #include +#include #include #include #include @@ -101,6 +102,11 @@ struct DynamicSparseAttentionCase { std::vector expected_output; std::vector expected_present_key; std::vector expected_present_value; + + // |actual - expected| <= absolute + relative * |expected| for the output; a negative relative tolerance keeps the + // default of the output type. + float output_abs_tolerance = 0.005f; + float output_rel_tolerance = -1.0f; }; template @@ -266,7 +272,7 @@ void RunDynamicSparseAttentionCase( if (c.expected_validation_status.has_value()) { tester.AddOutput("validation_status", {1}, {*c.expected_validation_status}); } - tester.SetOutputTolerance(0.005f); + tester.SetOutputTolerance(c.output_abs_tolerance, c.output_rel_tolerance); std::vector> execution_providers; execution_providers.push_back(std::move(execution_provider)); @@ -604,32 +610,50 @@ TEST(DynamicSparseAttentionTest, GroupedDecodeRatio12ReferenceParity_CUDA) { } } -// Selected-only attention without a softmax sink, the configuration the tensor-core kernels serve. With -// invalid_entries every fifth selected index is replaced by one that must be skipped (negative, in the future of the -// query row, or outside the cache). +// Selected-only attention without a softmax sink, the configuration the tensor-core kernels serve. The inputs are first +// rounded the way the operator receives them (the query and a floating-point cache as BFloat16; the codes of an INT8 +// cache are exact, its scales stay float), then the scores, the max-subtracted weights, the numerator and the +// denominator are evaluated in double precision, so that neither a large common logit nor a tiny weight overflows or +// underflows the reference. With invalid_entries every fifth selected index is replaced by one that must be skipped +// (negative, in the future of the query row, or outside the cache). void ComputeSelectedOnlyPrefillExpectedOutput(DynamicSparseAttentionCase& c, bool int8_cache) { + const auto bf16 = [](float value) { return static_cast(static_cast(BFloat16(value))); }; + std::vector scores; + std::vector offsets; + std::vector numerator(c.head_size); for (int row = 0; row < c.batch_size * c.sequence_length; ++row) { const int64_t b = row / c.sequence_length; const int64_t position = c.seqlens_k[b] + 1 - c.sequence_length + row % c.sequence_length; for (int h = 0; h < c.num_heads; ++h) { const int64_t kv = h / (c.num_heads / c.kv_num_heads); const int64_t out = (row * c.num_heads + h) * c.head_size; - std::vector numerator(c.head_size, 0.0f); - double denominator = 0.0; + scores.clear(); + offsets.clear(); for (int i = 0; i < c.selected_counts[row]; ++i) { const int index = c.selected_indices[row * c.max_selected + i]; if (index < 0 || index > position || index >= c.cache_sequence_length) continue; const int64_t offset = ((b * c.kv_num_heads + kv) * c.cache_sequence_length + index) * c.head_size; double dot = 0.0; for (int d = 0; d < c.head_size; ++d) { - const float key_scale = int8_cache ? c.key_scale[kv * c.head_size + d] : 1.0f; - dot += c.query[out + d] * c.expected_present_key[offset + d] * key_scale; + const double key = int8_cache ? static_cast(c.expected_present_key[offset + d]) * + static_cast(c.key_scale[kv * c.head_size + d]) + : bf16(c.expected_present_key[offset + d]); + dot += bf16(c.query[out + d]) * key; } - const double weight = std::exp(dot * c.scale); + scores.push_back(dot * static_cast(c.scale)); + offsets.push_back(offset); + } + const double max_score = scores.empty() ? 0.0 : *std::max_element(scores.begin(), scores.end()); + double denominator = 0.0; + std::fill(numerator.begin(), numerator.end(), 0.0); + for (size_t i = 0; i < scores.size(); ++i) { + const double weight = std::exp(scores[i] - max_score); denominator += weight; for (int d = 0; d < c.head_size; ++d) { - const float value_scale = int8_cache ? c.value_scale[kv * c.head_size + d] : 1.0f; - numerator[d] += static_cast(weight * c.expected_present_value[offset + d] * value_scale); + const double value = int8_cache ? static_cast(c.expected_present_value[offsets[i] + d]) * + static_cast(c.value_scale[kv * c.head_size + d]) + : bf16(c.expected_present_value[offsets[i] + d]); + numerator[d] += weight * value; } } for (int d = 0; d < c.head_size; ++d) { @@ -811,9 +835,10 @@ TEST(DynamicSparseAttentionTest, Int8PrefillQueryBeyondHalfRange_CUDA) { OpTester::ExpectResult::kExpectSuccess, "", configs); } -// Two selected candidates (cache positions 0 and 1) of an otherwise empty INT8 cache, so that the scores and the output -// follow from a few numbers: every query head holds `query` on channel 0, the keys hold key_codes on channel 0 and the -// values hold value_codes on every channel. The first row of the case selects nothing and outputs zeros. +// Two selected candidates of an otherwise empty INT8 cache, so that the scores and the output follow from a few numbers: +// every query head holds `query` on channel 0, the keys hold key_codes on channel 0 and the values hold value_codes on +// every channel. Candidate 0 is at cache position 0 and candidate 1 at second_position (negative: not selected). The +// first row of the case selects nothing and outputs zeros. struct TwoCandidateInt8Spec { float query; float key_scale; @@ -821,11 +846,43 @@ struct TwoCandidateInt8Spec { float scale; float key_codes[2]; float value_codes[2]; + int second_position = 1; + // Positive: every selecting row lists this many entries, candidate 0 first, candidate 1 last and -1 in between, which + // spreads the candidates over 64-candidate tiles (and over decode splits). Otherwise the row lists the candidates + // ahead of the -1 padding of the default 512 entries. + int list_length = 0; + // Checks the output within one bfloat16 ulp rather than the default 0.005 + 1% tolerance, for outputs that are tiny + // or whose wrong result differs from the right one by less than the default tolerance. + bool exact = false; + // Replaces the single-channel pattern (query, key codes and scales) before the reference is computed; receives + // second_position. + std::function customize; }; +// Sets one INT8 key code of every batch item and KV head. +void SetInt8KeyCode(DynamicSparseAttentionCase& c, int position, int channel, float code) { + for (int b = 0; b < c.batch_size; ++b) { + for (int kv = 0; kv < c.kv_num_heads; ++kv) { + c.past_key[((b * c.kv_num_heads + kv) * c.cache_sequence_length + position) * c.head_size + channel] = code; + } + } +} + DynamicSparseAttentionCase MakeTwoCandidateInt8Case(const TwoCandidateInt8Spec& spec, int num_heads, int sequence_length) { auto c = MakeSelectedOnlyPrefillCase(true, num_heads, 2, false, sequence_length); + const int rows = static_cast(c.batch_size * c.sequence_length); + const bool has_second = spec.second_position >= 0; + if (spec.list_length > 0) c.max_selected = spec.list_length; + c.selected_indices.assign(static_cast(rows * c.max_selected), -1); + for (int row = 0; row < rows; ++row) { + c.selected_counts[row] = row == 0 ? 0 : spec.list_length > 0 ? spec.list_length + : (has_second ? 2 : 1); + if (row == 0) continue; + int32_t* entries = c.selected_indices.data() + static_cast(row * c.max_selected); + entries[0] = 0; + if (has_second) entries[spec.list_length > 0 ? spec.list_length - 1 : 1] = spec.second_position; + } c.scale = spec.scale; std::fill(c.query.begin(), c.query.end(), 0.0f); std::fill(c.key_scale.begin(), c.key_scale.end(), spec.key_scale); @@ -834,41 +891,79 @@ DynamicSparseAttentionCase MakeTwoCandidateInt8Case(const TwoCandidateInt8Spec& std::fill(c.value.begin(), c.value.end(), 0.0f); std::fill(c.past_key.begin(), c.past_key.end(), 0.0f); std::fill(c.past_value.begin(), c.past_value.end(), 0.0f); + const int positions[2] = {0, spec.second_position}; for (int b = 0; b < c.batch_size; ++b) { for (int kv = 0; kv < c.kv_num_heads; ++kv) { const int64_t offset = (b * c.kv_num_heads + kv) * c.cache_sequence_length * c.head_size; for (int i = 0; i < 2; ++i) { - c.past_key[offset + i * c.head_size] = spec.key_codes[i]; - std::fill_n(c.past_value.begin() + offset + i * c.head_size, c.head_size, spec.value_codes[i]); + if (positions[i] < 0) continue; + c.past_key[offset + positions[i] * c.head_size] = spec.key_codes[i]; + std::fill_n(c.past_value.begin() + offset + positions[i] * c.head_size, c.head_size, spec.value_codes[i]); } } } + for (size_t i = 0; i < c.query.size(); i += c.head_size) c.query[i] = spec.query; + if (spec.customize) spec.customize(c, spec.second_position); c.expected_present_key = c.past_key; c.expected_present_value = c.past_value; - for (size_t i = 0; i < c.query.size(); i += c.head_size) c.query[i] = spec.query; ComputeSelectedOnlyPrefillExpectedOutput(c, true); + if (spec.exact) { + c.output_abs_tolerance = 0.0f; + c.output_rel_tolerance = 0x1p-7f; + } return c; } -// Every row but the first outputs `expected` on every channel. Prefill (9 rows per batch) and decode (3 rows) run -// with the tensor-core INT8 kernels both disabled (reference fallback) and enabled. -void RunTwoCandidateInt8Case(const TwoCandidateInt8Spec& spec, float expected) { - for (bool decode : {false, true}) { - const auto c = MakeTwoCandidateInt8Case(spec, 12, decode ? 3 : 9); +// What RunTwoCandidateInt8Case runs: prefill (9 rows per batch item) and decode with each of the sequence lengths. +// The tensor-core decode kernel also runs with each tiles-per-split setting (the other kernels ignore it). +struct TwoCandidateInt8Run { + bool prefill = true; + std::vector decode_sequence_lengths = {3}; + std::vector tiles_per_split = {"0"}; +}; + +// Every row but the first outputs `expected` on every channel. The tensor-core INT8 kernels run both disabled +// (reference fallback) and enabled. +void RunTwoCandidateInt8Case(const TwoCandidateInt8Spec& spec, float expected, const TwoCandidateInt8Run& run = {}) { + std::vector> shapes; // (decode, sequence length) + if (run.prefill) shapes.emplace_back(false, 9); + for (int sequence_length : run.decode_sequence_lengths) shapes.emplace_back(true, sequence_length); + for (const auto& [decode, sequence_length] : shapes) { + const auto c = MakeTwoCandidateInt8Case(spec, 12, sequence_length); for (size_t i = 0; i < c.expected_output.size(); i += c.head_size) { const size_t row = i / (c.num_heads * c.head_size); - ASSERT_NEAR(c.expected_output[i], row == 0 ? 0.0f : expected, 1e-6f * std::max(1.0f, std::fabs(expected))); + ASSERT_NEAR(c.expected_output[i], row == 0 ? 0.0f : expected, 1e-6 * std::fabs(expected)); } for (const char* int8_tc : {"0", "1"}) { - SCOPED_TRACE(std::string(decode ? "decode" : "prefill") + ", int8_tc=" + int8_tc); - RunDynamicSparseAttentionCase( - c, DefaultCudaExecutionProvider(), OpTester::ExpectResult::kExpectSuccess, "", - decode ? DecodeKernelConfigs(int8_tc) : PrefillKernelConfigs("1", int8_tc)); - if (::testing::Test::HasFailure()) return; + const std::vector tiles_per_split_values = + decode && std::string(int8_tc) == "1" ? run.tiles_per_split : std::vector{"0"}; + for (const char* tiles_per_split : tiles_per_split_values) { + SCOPED_TRACE(std::string(decode ? "decode" : "prefill") + ", sequence_length=" + + std::to_string(sequence_length) + ", int8_tc=" + int8_tc + + ", tiles_per_split=" + tiles_per_split); + RunDynamicSparseAttentionCase( + c, DefaultCudaExecutionProvider(), OpTester::ExpectResult::kExpectSuccess, "", + decode ? DecodeKernelConfigs(int8_tc, tiles_per_split) : PrefillKernelConfigs("1", int8_tc)); + if (::testing::Test::HasFailure()) return; + } } } } +// Runs the case with its candidates in one 64-candidate tile and again with a list of 129 entries that puts the +// candidates in the first and the last of three tiles, which the tensor-core decode kernel spreads over up to three +// splits (two of them active) and the prefill kernel consumes in one CTA. +void RunTwoCandidateInt8CaseOverTiles(TwoCandidateInt8Spec spec, float expected, bool prefill = true) { + TwoCandidateInt8Run run; + run.prefill = prefill; + RunTwoCandidateInt8Case(spec, expected, run); + if (::testing::Test::HasFailure()) return; + spec.list_length = 129; + if (spec.second_position >= 0) spec.second_position = 128; + run.tiles_per_split = {"0", "1", "2", "3"}; + RunTwoCandidateInt8Case(spec, expected, run); +} + // The tensor-core INT8 kernels round Q * key_scale to fp16, which flushes products below 2^-25 to zero; a softmax scale // above 1 would amplify that error, so those calls use the BF16-dequantizing prefill and CUDA-core decode kernels. // Every product here is 2^-30 (zero in fp16) and the scale of 2^30 turns it into scores of +-1, so the output is @@ -888,12 +983,14 @@ TEST(DynamicSparseAttentionTest, Int8TensorCoreLargeSoftmaxScale_CUDA) { } // The tensor-core INT8 kernels pass the softmax probabilities to the PV mma as fp16 holding 2^15 * p (the factor leaves -// in the epilogue), so a candidate keeps its contribution down to p = 2^-40 where plain fp16 flushes p below 2^-25; -// when value_scale > 2^16 (kInt8TcMaxValueScale), the CTA evaluates the row/split in fp32 (AttendRowFp32) so that -// probabilities below 2^-40 amplified by arbitrarily large finite value scales are preserved as well: +// in the epilogue), so a candidate keeps its contribution down to p = e^-20 (kInt8TcTailDepth) where plain fp16 flushes +// p below 2^-25. A CTA with a value scale above 2^16 (kInt8TcMaxValueScale) and a valid candidate more than +// kInt8TcTailDepth nats below its running maximum evaluates the row/split in fp32 (AttendRowFp32) instead, so that +// probabilities that leave the fp16 range, amplified by an arbitrarily large finite value scale, are preserved as well: // 1. query = 288 -> score -18, p = e^-18 ~= 2^-25.97 (flushes in unscaled fp16, kept by 2^15 * p on the TC path with -// value_scale = 2^16 <= kInt8TcMaxValueScale and value code 100); -// 2. query = 320 -> score -20, p = e^-20 ~= 2^-28.85, value_scale = 2^28; +// value_scale = 2^16 and value code 100); +// 2. query = 320 -> score -20, p = e^-20 ~= 2^-28.85, value_scale = 2^28: the depth equals kInt8TcTailDepth, so the +// tensor-core path still serves it; // 3. query = 480 -> score -30, p = e^-30 ~= 2^-43.28 (flushes even 2^15 * p in fp16), value_scale = 2^60, output // 2^60 * e^-30 / (1 + e^-30) ~= 107886.05 via the AttendRowFp32 fallback. TEST(DynamicSparseAttentionTest, Int8TensorCoreTinyProbabilityLargeValueScale_CUDA) { @@ -939,9 +1036,203 @@ TEST(DynamicSparseAttentionTest, Int8TensorCoreTinyProbabilityLargeValueScale_CU } } +// The cases below fix the outputs of the tensor-core INT8 kernels on numerical corners with exact arithmetic: the output +// follows from a few numbers, is checked to one bfloat16 ulp, and is also produced with the candidates in the first and +// the last of three 64-candidate tiles, which the decode kernel spreads over up to three splits (two of them active). + +// Q * key_scale is exact in fp32 here (32768 and -32512 * 129 / 128 = -32766 for a channel pair), and the all-ones key +// of candidate 0 makes the 128 pairs add up to 256, a logit gap of 16 at scale 1/16 against the all-zero key of +// candidate 1. -32766 is not an fp16 number though: rounding it to fp16 gives -32768, cancels every pair and turns the +// gap into 0 (output 0.5 instead of 1 - 1.1e-7). +TEST(DynamicSparseAttentionTest, Int8TensorCoreQueryFoldingRoundingError_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + TwoCandidateInt8Spec spec{}; + spec.key_scale = 1.0f; + spec.value_scale = 1.0f; + spec.scale = 0.0625f; + spec.value_codes[0] = 1.0f; + spec.exact = true; + spec.customize = [](DynamicSparseAttentionCase& c, int) { + for (size_t i = 0; i < c.query.size(); ++i) c.query[i] = i % 2 == 0 ? 32768.0f : -32512.0f; + for (size_t i = 0; i < c.key_scale.size(); ++i) c.key_scale[i] = i % 2 == 0 ? 1.0f : 1.0078125f; + for (int d = 0; d < c.head_size; ++d) SetInt8KeyCode(c, 0, d, 1.0f); + }; + RunTwoCandidateInt8CaseOverTiles(spec, 1.0f); +} + +// Logits of +-2^-13 and values of +-4064 (value code +-127, value scale 32) give 4064 * tanh(2^-13) = 127 / 256. The two +// probabilities, 0.5 +- 2^-14, both round to 0.5 in one fp16 tile (the spacing is 2^-11 above 0.5 and 2^-12 below), so +// the difference of the two values would cancel to 0; the tensor-core decode kernel passes P as a high and a low fp16 +// part to keep it. The prefill kernels round each probability to a single 16-bit tile for the PV mma, as the kernels +// before them did, so only decode is covered. +TEST(DynamicSparseAttentionTest, Int8TensorCoreDecodeSignedValueCancellation_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + TwoCandidateInt8Spec spec{}; + spec.query = 0x1p-9f; + spec.key_scale = 1.0f; + spec.value_scale = 32.0f; + spec.scale = 0.0625f; + spec.key_codes[0] = 1.0f; + spec.key_codes[1] = -1.0f; + spec.value_codes[0] = 127.0f; + spec.value_codes[1] = -127.0f; + spec.exact = true; + RunTwoCandidateInt8CaseOverTiles(spec, 127.0f / 256.0f, /*prefill=*/false); +} + +// Both candidates score 8290560 * scale (query 65280, close to the fp16 maximum, times key code 127 on channel 0), and +// candidate 1 adds 0.5 * scale from a second channel. Softmax weights computed in the log2 domain lose that gap to the +// float spacing of the common logit (0.0625 near 7.5e5 in log2 units, against a gap of 0.045 at scale 1/16) and give +// 0.5 instead of 1 / (1 + exp(-0.5 * scale)); the natural-log domain keeps it, since every logit here is an exact +// float. +TEST(DynamicSparseAttentionTest, Int8TensorCoreLargeCommonLogit_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + for (float scale : {1.0f, 0.0625f}) { + SCOPED_TRACE("scale=" + std::to_string(scale)); + TwoCandidateInt8Spec spec{}; + spec.query = 65280.0f; + spec.key_scale = 1.0f; + spec.value_scale = 1.0f; + spec.scale = scale; + spec.key_codes[0] = 127.0f; + spec.key_codes[1] = 127.0f; + spec.value_codes[1] = 1.0f; + spec.exact = true; + spec.customize = [](DynamicSparseAttentionCase& c, int second_position) { + for (size_t i = 1; i < c.query.size(); i += c.head_size) c.query[i] = 0.5f; + SetInt8KeyCode(c, second_position, 1, 1.0f); + }; + RunTwoCandidateInt8CaseOverTiles(spec, static_cast(1.0 / (1.0 + std::exp(-0.5 * scale)))); + if (::testing::Test::HasFailure()) return; + } +} + +// A single candidate with value code 127 and value scale 2^-135 outputs 127 * 2^-135 = 2.9e-39, a subnormal in float +// that rounds to the bfloat16 subnormal 2^-128; a kernel that folds the value scale into an fp16 operand or multiplies +// it in after an fp16 conversion flushes it to 0. The second layout lists the candidate first among 129 entries, which +// leaves two splits empty with three tiles per row. +TEST(DynamicSparseAttentionTest, Int8TensorCoreTinyValueScale_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + TwoCandidateInt8Spec spec{}; + spec.key_scale = 1.0f; + spec.value_scale = 0x1p-135f; + spec.scale = 0.0625f; + spec.value_codes[0] = 127.0f; + spec.second_position = -1; + spec.exact = true; + RunTwoCandidateInt8CaseOverTiles(spec, 127.0f * 0x1p-135f); +} + +// The tensor-core decode kernel runs at most 32 splits per row (and no more than the workspace holds, 16 for the eight +// rows of a batch item here) however many 64-candidate tiles the list has, so a tiles-per-split setting is a target: +// 8192 entries are 128 tiles, and a setting of 1 requests 128 splits but gets 32 splits of four tiles. Whatever the +// split count, every candidate must reach the merge: the logits are all equal and the values are 2 and 6, so a split +// that is dropped or counted twice gives 2, 6 or something else instead of 4. The lists are 2048 entries (32 tiles, +// exactly the cap for 1 tile per split), 2112 (33 tiles, one over) and 8192, with the second candidate in the last tile. +TEST(DynamicSparseAttentionTest, Int8TensorCoreDecodeTilesPerSplitCaps_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + for (int list_length : {2048, 2112, 8192}) { + SCOPED_TRACE("list_length=" + std::to_string(list_length)); + TwoCandidateInt8Spec spec{}; + spec.key_scale = 1.0f; + spec.value_scale = 1.0f; + spec.scale = 0.0625f; + spec.value_codes[0] = 2.0f; + spec.value_codes[1] = 6.0f; + spec.second_position = 300; + spec.list_length = list_length; + TwoCandidateInt8Run run; + run.prefill = false; + run.decode_sequence_lengths = {1, 8}; + run.tiles_per_split = {"0", "1", "2", "3", "64"}; + RunTwoCandidateInt8Case(spec, 4.0f, run); + if (::testing::Test::HasFailure()) return; + } +} + +// Runs a case built for a prefill of 9 rows per batch item (decode == false) or a decode of 3 with the tensor-core +// kernels; the decode kernel gets the automatic, 1 and 3 tiles per split. The BF16-dequantizing prefill kernel is left +// out: it rounds the dequantized keys and values to bfloat16, which the reference of these cases does not allow for. +void RunInt8TensorCoreCase(const DynamicSparseAttentionCase& c, bool decode) { + const std::vector tiles_per_split_values = + decode ? std::vector{"0", "1", "3"} : std::vector{"0"}; + for (const char* tiles_per_split : tiles_per_split_values) { + SCOPED_TRACE(std::string(decode ? "decode" : "prefill") + ", tiles_per_split=" + tiles_per_split); + RunDynamicSparseAttentionCase( + c, DefaultCudaExecutionProvider(), OpTester::ExpectResult::kExpectSuccess, "", + decode ? DecodeKernelConfigs("1", tiles_per_split) : PrefillKernelConfigs("1", "1")); + if (::testing::Test::HasFailure()) return; + } +} + +// Random rows with value scales between 2^14 and 2^17 (outputs up to 10^7), so that the tail-depth guard of the +// tensor-core kernels sees value scales above 2^16 on random data. Queries scaled by 1, 16 and 32 spread the scores of a +// head over about 0.7, 11 and 22 nats: the first two stay on the tensor-core path, and some CTAs of the third find a +// candidate more than 20 nats below the maximum and take the fp32 fallback. The tolerance is the 0.005 of the unscaled +// fixture scaled with the values. +TEST(DynamicSparseAttentionTest, Int8TensorCoreLargeValueScaleRandom_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + for (float query_multiplier : {1.0f, 16.0f, 32.0f}) { + for (int sequence_length : {9, 3}) { + SCOPED_TRACE("query_multiplier=" + std::to_string(query_multiplier) + + ", sequence_length=" + std::to_string(sequence_length)); + auto c = MakeSelectedOnlyPrefillCase(true, 24, 509, true, sequence_length); + for (float& value_scale : c.value_scale) value_scale *= 0x1p22f; + for (float& value : c.value) value *= 0x1p22f; + for (float& query : c.query) query *= query_multiplier; + ComputeSelectedOnlyPrefillExpectedOutput(c, true); + c.output_abs_tolerance = 0.005f * 0x1p22f; + RunInt8TensorCoreCase(c, sequence_length == 3); + if (::testing::Test::HasFailure()) return; + } + } +} + +// Random INT8 rows whose folded query products Q * key_scale are not fp16 numbers: key scales are random in +// [0.01, 0.03] and queries are drawn from N(0, 0.25^2) (rounded to bfloat16 as the kernels receive them). The other random +// fixtures fold to exact fp16 products, so they never exercise the rounding of the folded query. A query multiplier +// of 1 keeps the accumulated rounding error of a head far below the error budget of the fp16 scores, and 16 pushes it over, +// into the fp32 scores. +TEST(DynamicSparseAttentionTest, Int8TensorCoreRandomQueryFolding_CUDA) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; + std::mt19937 gen(29); + std::uniform_real_distribution key_scale_distribution(0.01f, 0.03f); + std::normal_distribution query_distribution(0.0f, 0.25f); + for (float query_multiplier : {1.0f, 16.0f}) { + for (int num_heads : {12, 24}) { + for (int selected_count : {65, 509}) { + for (int sequence_length : {9, 3}) { + SCOPED_TRACE("query_multiplier=" + std::to_string(query_multiplier) + + ", num_heads=" + std::to_string(num_heads) + + ", selected_count=" + std::to_string(selected_count) + + ", sequence_length=" + std::to_string(sequence_length)); + auto c = MakeSelectedOnlyPrefillCase(true, num_heads, selected_count, true, sequence_length); + const size_t channels = static_cast(c.kv_num_heads * c.head_size); + for (float& key_scale : c.key_scale) key_scale = key_scale_distribution(gen); + // The new tokens are the codes of the fixture times the new scales; expected_present_key keeps the codes. + for (size_t i = 0; i < c.key.size(); ++i) { + c.key[i] = static_cast(static_cast(i % 11) - 5) * c.key_scale[i % channels]; + } + for (float& query : c.query) query = query_multiplier * query_distribution(gen); + ComputeSelectedOnlyPrefillExpectedOutput(c, true); + RunInt8TensorCoreCase(c, sequence_length == 3); + if (::testing::Test::HasFailure()) return; + } + } + } + } +} + // Explicit ep.cuda.* entries take precedence over the environment, which is consulted only when no entry is set; an // invalid value from either source fails kernel construction (with an exception, so those cases are skipped in builds -// without exceptions). +// without exceptions). The switches accept exactly the strings "0" and "1". TEST(DynamicSparseAttentionTest, PrefillSessionOptions_CUDA) { if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA EP not available."; if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16."; @@ -975,11 +1266,19 @@ TEST(DynamicSparseAttentionTest, PrefillSessionOptions_CUDA) { RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_success); } #ifndef ORT_NO_EXCEPTIONS - for (const char* invalid : {"2", "-1", "", "1x", "true"}) { - SCOPED_TRACE(std::string("value='") + invalid + "'"); + // Only the strings "0" and "1" are switch values, not other spellings of the same numbers. + for (const char* key : {kOrtSessionOptionsCudaDsaPrefillKernel, kOrtSessionOptionsCudaDsaPrefillInt8Tc}) { + for (const char* invalid : {"2", "-1", "", "1x", "true", "00", "01", "-0", "000001", "+1", " 1", "1 "}) { + SCOPED_TRACE(std::string(key) + "='" + invalid + "'"); + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_failure, key, + {{key, invalid}}); + } + } + for (const char* invalid : {"01", "-0", "+1"}) { + SCOPED_TRACE(std::string("environment value='") + invalid + "'"); + ScopedEnvironmentVariables env{EnvVarMap{{"ORT_DSA_PREFILL_INT8_TC", invalid}}}; RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_failure, - kOrtSessionOptionsCudaDsaPrefillInt8Tc, - {{kOrtSessionOptionsCudaDsaPrefillInt8Tc, invalid}}); + "ORT_DSA_PREFILL_INT8_TC"); } #endif } @@ -1108,12 +1407,20 @@ TEST(DynamicSparseAttentionTest, DecodeSessionOptions_CUDA) { RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_success); } #ifndef ORT_NO_EXCEPTIONS - for (const char* invalid : {"2", "-1", "", "1x", "true"}) { + // Only the strings "0" and "1" are switch values, not other spellings of the same numbers. + for (const char* invalid : {"2", "-1", "", "1x", "true", "00", "01", "-0", "000001", "+1", " 1", "1 "}) { SCOPED_TRACE(std::string("value='") + invalid + "'"); RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_failure, kOrtSessionOptionsCudaDsaDecodeInt8Tc, {{kOrtSessionOptionsCudaDsaDecodeInt8Tc, invalid}}); } + for (const char* invalid : {"01", "-0", "+1"}) { + SCOPED_TRACE(std::string("environment value='") + invalid + "'"); + ScopedEnvironmentVariables env{EnvVarMap{{"ORT_DSA_DECODE_INT8_TC", invalid}}}; + RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_failure, + "ORT_DSA_DECODE_INT8_TC"); + } + // The tiles-per-split threshold is an integer, not a switch. for (const char* invalid : {"-1", "", "1x", "many"}) { SCOPED_TRACE(std::string("value='") + invalid + "'"); RunDynamicSparseAttentionCase(c, DefaultCudaExecutionProvider(), expect_failure, From 2dedd0c4ff26ffaa1f2bc2b51e061bd1826c9424 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 31 Oct 2026 20:43:18 +0000 Subject: [PATCH 12/14] [CUDA] DynamicSparseAttention: document numerics, fallbacks and rollout of the INT8 tensor-core kernels, add an op-level benchmark docs/contrib_ops/cuda/dynamic_sparse_attention.md: - Dispatch switches: both INT8 tensor-core switches are opt-in (default 0) and accept exactly "0" and "1". The tiles-per-split rule (smallest power of two that fits one wave of resident CTAs, at most 32 splits) is described with its measured behaviour: the time against tiles per split on an H200, the regret against the best swept value over 20 shapes (mean 1.8%, at most 14.8%), the emulated one and two CTAs per SM, and the A100 split count that follows from the rule (derived, not measured). - Numerics and fallbacks: when a CTA takes the FP32 path (Q' outside FP16 or non-finite, eps above the budget, non-finite value_scale, a candidate more than 20 nats below the running maximum with abs(value_scale) > 2^16), the eps against query magnitude table, the measured cost of the FP32 path (up to 21x slower than the previous kernels when every CTA takes it, accuracy kept) and the tests that cover each condition. - Rollout: the kernels stay opt-in until the final revision is qualified end to end; the earlier accuracy result and the end-to-end timings are labelled as results for an earlier revision. - Benchmark: operator-level decode and prefill tables against the kernels that the tensor-core kernels replace, with the RMS error against FP64. benchmark_dynamic_sparse_attention.py: single-node DynamicSparseAttention models with the shapes of the Qwen3.8-Flash-Next layers (or the scales of a model's node), one session per kernel configuration, the GPU kernel time from the CUPTI profiler, the median of paired per-round ratios with alternating order, the RMS error against an FP64 reference, and options that reproduce the FP32 fallbacks (--query-std, --value-scale-multiplier, --key-scale-multiplier, ...). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cuda/dynamic_sparse_attention.md | 214 +++++++-- .../benchmark_dynamic_sparse_attention.py | 422 ++++++++++++++++++ 2 files changed, 604 insertions(+), 32 deletions(-) create mode 100644 onnxruntime/test/python/transformers/benchmark_dynamic_sparse_attention.py diff --git a/docs/contrib_ops/cuda/dynamic_sparse_attention.md b/docs/contrib_ops/cuda/dynamic_sparse_attention.md index ed96cf3c727a7..337dd5e4d9a87 100644 --- a/docs/contrib_ops/cuda/dynamic_sparse_attention.md +++ b/docs/contrib_ops/cuda/dynamic_sparse_attention.md @@ -47,11 +47,14 @@ the code is only reported through `validation_status`. `ORT_DSA_INT8_VECTOR_STAGING` and `ORT_DSA_SHARED_SCALE_MERGE` are read once per process (cached in function-local statics) when the grouped decode path is first launched, so set them before the first operator invocation. Only `0` and `1` are accepted. Neither changes schemas or cache layouts. -`ORT_DSA_PREFILL_KERNEL`, `ORT_DSA_PREFILL_INT8_TC` and `ORT_DSA_DECODE_INT8_TC` (all default `1`) are read once -at kernel construction. Session options `ep.cuda.dsa_prefill_kernel`, `ep.cuda.dsa_prefill_int8_tc` and -`ep.cuda.dsa_decode_int8_tc` take precedence over their corresponding environment variables. -They accept only `0` or `1`; changing a value requires a new session. The loaded prefill and decode -specializations must have a PTX target of SM80+; otherwise dispatch uses the generic fallback. +`ORT_DSA_PREFILL_KERNEL`, `ORT_DSA_PREFILL_INT8_TC` and `ORT_DSA_DECODE_INT8_TC` are read once at kernel +construction. `ORT_DSA_PREFILL_KERNEL` defaults to `1`. The two INT8 tensor-core switches default to `0`: they are +opt-in, and without them the operator runs the BF16-dequantizing prefill and the CUDA-core grouped decode, bit for +bit as before (see [Rollout](#rollout)). Session options `ep.cuda.dsa_prefill_kernel`, `ep.cuda.dsa_prefill_int8_tc` +and `ep.cuda.dsa_decode_int8_tc` take precedence over their corresponding environment variables. Each of them +accepts exactly the one-character string `0` or `1` (no whitespace, sign or leading zeros: `01`, `-0` and `true` are +rejected); changing a value requires a new session. The loaded prefill and decode specializations must have a PTX +target of SM80+; otherwise dispatch uses the generic fallback. `ORT_DSA_DECODE_INT8_TC_TILES_PER_SPLIT` (session option `ep.cuda.dsa_decode_int8_tc_tiles_per_split`) is a tuning and test knob with the same precedence: an integer in `[0, INT_MAX]`, `0` keeps the automatic split heuristic. Invalid values of any of these options fail session initialization. @@ -70,18 +73,123 @@ and 8 otherwise. | | | | `0` | `MergeGroupedDecodeDynamicSparseAttentionKernel`: per-channel `expf` over all splits. | | `ORT_DSA_PREFILL_KERNEL` | `1` | BF16, not grouped decode, `sequence_length >= 8`, `head_size == 256`, `selected_only` from the main cache, `smooth_softmax = 0`, SM80+ with enough opt-in shared memory | `1` | `DynamicSparseAttentionPrefillKernel` tensor-core prefill. | | | | | `0` | Generic fused/split or grouped prefill kernels. | -| `ORT_DSA_PREFILL_INT8_TC` | `1` | The `ORT_DSA_PREFILL_KERNEL` conditions with an INT8 main cache, 16-byte-aligned `present_key` and `present_value`, an 8-byte-aligned `output`, a softmax `scale <= 1` and enough opt-in shared memory for `dsa_prefill::kInt8SharedBytes` | `1` | `DynamicSparseAttentionInt8TcKernel` ([dynamic_sparse_attention_prefill_int8.cuh](../../../onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh)): the gathered K/V rows stay INT8 in shared memory and are converted exactly to FP16 in registers for FP16 `mma` with FP32 accumulation. `key_scale` is folded into the query once per block and `value_scale` is applied to the output once. Rounding `Q * key_scale` to FP16 flushes products below `2^-25`, which moves a score by at most `2^-10 * scale`, so a call with `scale > 1` (the default is `1 / sqrt(head_size)`) uses the kernel below. The FP16 probability tile holds `2^15 * p`, which the epilogue removes, so that a candidate keeps its contribution down to `p = 2^-40` (plain FP16 flushes below `2^-25`); the row sum stays unscaled FP32. A block whose scaled query does not fit FP16 (`abs(Q * key_scale) > 65504`, inf or NaN) or whose value scale exceeds `kInt8TcMaxValueScale` (`2^16`, above which probabilities flushed below `2^-40` can become significant after `value_scale`) is evaluated by an FP32 path (`AttendRowFp32`) instead of being clipped, which would break the cancellation between large terms. That path is about 48x slower for the affected blocks (8192 rows on an H200: 2.75 ms to 131 ms); set `ep.cuda.dsa_prefill_int8_tc=0` for models whose scaled queries exceed the FP16 range, because the kernel below has no such limit. | +| `ORT_DSA_PREFILL_INT8_TC` | `0` | The `ORT_DSA_PREFILL_KERNEL` conditions with an INT8 main cache, 16-byte-aligned `present_key` and `present_value`, an 8-byte-aligned `output`, a softmax `scale <= 1` and enough opt-in shared memory for `dsa_prefill::kInt8SharedBytes` | `1` | `DynamicSparseAttentionInt8TcKernel` ([dynamic_sparse_attention_prefill_int8.cuh](../../../onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill_int8.cuh)): the gathered K/V rows stay INT8 in shared memory and are converted exactly to FP16 in registers for FP16 `mma` with FP32 accumulation. `key_scale` is folded into the query once per block (`Q' = Q * key_scale`, rounded once to FP16) and `value_scale` is applied to the output once. Scores and the running maximum are kept in the natural-log domain. A softmax `scale > 1` (the default is `1 / sqrt(head_size)`) uses the kernel below, because the rounding bound grows with it. The FP16 probability tile holds `2^15 * p`, so a candidate keeps its contribution down to `p = 2^-40` (plain FP16 flushes below `2^-25`); the row sum stays unscaled FP32. Prefill keeps a single FP16 probability tile, as the BF16-dequantizing kernel does with a coarser BF16 one. A block whose query or value scales leave the validated range is evaluated by an FP32 path (`AttendRowFp32`) instead of being clipped or rounded, which would break the cancellation between large terms; the conditions, their cost and how likely they are are described under [Numerics and fallbacks](#numerics-and-fallbacks). | | | | | `0` | `DynamicSparseAttentionPrefillKernel`: dequantizes every gathered tile to BF16 in shared memory. | -| `ORT_DSA_DECODE_INT8_TC` | `1` | Grouped decode (conditions above) with BF16, an INT8 main cache, `head_size == 256`, `selected_only` from the main cache (no `local_plus_selected`), `smooth_softmax = 0`, a softmax `scale <= 1`, at most `kMaxGroup` query heads per KV head, 16-byte-aligned `prepared_query`, `present_key`, `present_value`, `key_scale`, `value_scale` and attention workspace, SM80+ and enough opt-in shared memory for `dsa_prefill::kInt8SharedBytes` | `1` | `DynamicSparseAttentionInt8TcKernel`: the tensor-core pipeline of the row above, split over the 64-candidate tiles of each (row, KV head). Each CTA of the (row, KV head, split) grid writes an un-normalized partial output, the natural-log row maximum and the row sum in the layout of the grouped decode merge kernel, which is reused unchanged; the `2^15` probability factor of the row above is removed before a partial is written. The split count targets 384 CTAs (one wave of three per SM, at most 32 splits). A CTA whose scaled query does not fit FP16 evaluates its candidates in FP32 (`AttendRowFp32`) and publishes an FP32 partial in the same layout, so the range-safe fallback also holds for the splits. | +| `ORT_DSA_DECODE_INT8_TC` | `0` | Grouped decode (conditions above) with BF16, an INT8 main cache, `head_size == 256`, `selected_only` from the main cache (no `local_plus_selected`), `smooth_softmax = 0`, a softmax `scale <= 1`, at most `kMaxGroup` query heads per KV head, 16-byte-aligned `prepared_query`, `present_key`, `present_value`, `key_scale`, `value_scale` and attention workspace, SM80+ and enough opt-in shared memory for `dsa_prefill::kInt8SplitKSharedBytes` (75,264 B) | `1` | `DynamicSparseAttentionInt8TcKernel`: the tensor-core pipeline of the row above, split over the 64-candidate tiles of each (row, KV head). Each CTA of the (row, KV head, split) grid writes an un-normalized partial output, the natural-log row maximum and the row sum in the layout of the grouped decode merge kernel, which is reused unchanged. The probability tile is an FP16 high part plus an FP16 residual (about 22 significant bits) and both multiply the converted V fragments, because the signed sum of large value codes can cancel to a value far below the rounding of one FP16 probability. A partial is written as `(acc * 2^-15) * value_scale`. The split count follows `ORT_DSA_DECODE_INT8_TC_TILES_PER_SPLIT` below. A CTA that meets one of the conditions of [Numerics and fallbacks](#numerics-and-fallbacks) evaluates its candidates in FP32 (`AttendRowFp32`) and publishes an FP32 partial in the same layout, so the fallback also holds for the splits. | | | | | `0` | The CUDA-core `GroupedDecodeDynamicSparseAttentionKernel`, bit for bit as before. | -| `ORT_DSA_DECODE_INT8_TC_TILES_PER_SPLIT` | `0` | `ORT_DSA_DECODE_INT8_TC=1` and the conditions above | `0` | Automatic: `max(1, ceil(tiles * rows * kv_heads / 384))` tiles per split, with `rows = batch_size * sequence_length`; the split count never exceeds 32 or the split count the grouped decode workspace was sized for. | -| | | | `N > 0` | `N` 64-candidate tiles per split (the sweep around the automatic value was flat; used by the tests to force many or few splits). | +| `ORT_DSA_DECODE_INT8_TC_TILES_PER_SPLIT` | `0` | `ORT_DSA_DECODE_INT8_TC=1` and the conditions above | `0` | Automatic: the smallest power-of-two number of 64-candidate tiles per split for which all `ceil(max_selected / 64) * rows * kv_heads` tiles fit one wave of resident CTAs (`rows = batch_size * sequence_length`; resident CTAs = SM count times the CTAs per SM that the kernel's shared-memory and register footprint allow, queried with `cudaOccupancyMaxActiveBlocksPerMultiprocessor`: 396 on an H200). The split count is `ceil(tiles / tiles_per_split)`, at most 32 and never more than the grouped decode workspace was sized for. | +| | | | `N > 0` | A requested `N` tiles per split instead of the automatic choice. The 32-split cap, the workspace size and the even distribution of the tiles over the splits still apply, so the number of tiles of a split can exceed `N` (8192 candidates are 128 tiles: `N = 1` asks for 128 splits and runs 32 splits of 4 tiles). Used by the tests to force many or few splits. | + +The decode time is not flat in the tiles per split: with too few splits the CTAs are long and SMs stay idle, with too +many the grid spills into a second, partly filled wave. On an H200, 8 rows and 2048 candidates (16 pairs of row and KV +head with 32 tiles each), kernel plus merge, +`benchmark_dynamic_sparse_attention.py --rows 8 --tiles-per-split 0 1 2 4 8 16`: + +| Tiles per split | 1 | 2 (automatic) | 4 | 8 | 16 | +|---|---|---|---|---|---| +| Splits, CTAs | 32, 512 | 16, 256 | 8, 128 | 4, 64 | 2, 32 | +| Time | 18.8 us | 13.2 us | 12.7 us | 17.7 us | 29.2 us | + +The automatic rule is a heuristic, not the optimum. Over 20 shapes on the H200 (1, 2, 4 and 8 rows, 512 to 8192 +candidates, context 8192, each swept over 1, 2, 3, 4, 6, 8 and 16 tiles per split) it is on average 1.8% and at most +14.8% slower than the best swept value (4 rows and 2048 candidates: 11.9 us against 10.4 us at 2 tiles per split). The +previous target of about 384 CTAs is 2.7% and 14.2% slower on the same points: with 396 resident CTAs the two rules mostly +agree on an H200, and the larger shapes gain from the power-of-two choice (8 rows, 8192 candidates: 0.2% against 6.7%). +The rule matters when fewer CTAs are resident. With the kernel's shared memory padded on the H200 so that one or two CTAs +fit an SM (132 or 264 resident CTAs; an ad hoc harness that is not part of this change; 1 to 32 rows, 2048 candidates) +the rule is on average 0.4% (at most 2.4%) and 3.1% (13.4%) slower than the best swept value, where the 384-CTA target is +23.8% (44.9%) and 14.4% (38.4%) slower. A device with 108 SMs and two resident CTAs per SM (the 75,264 B of shared memory +allow two on an A100) is expected to choose 4 tiles per split, 8 splits, for 8 rows and 2048 candidates, the faster +setting in an A100 measurement reported for this change (22.1 us against 25.9 us at 16 splits). The A100 has not been +measured with this document's script; `ORT_DSA_DECODE_INT8_TC_TILES_PER_SPLIT` tunes other devices. When `split_count > 32` (floating-point caches only) the generic `MergeDynamicSparseAttentionKernel` is used regardless of `ORT_DSA_SHARED_SCALE_MERGE`. The tests `Int8CacheGroupedDecodeAlternateSwitchesInFreshProcess_CUDA` re-run the INT8 grouped-decode and merge cases in fresh processes with each non-default combination. +### Numerics and fallbacks + +The INT8 tensor-core kernels multiply `Q' = fp16(Q * key_scale)` with the exactly converted INT8 K rows and the FP16 +probability tile with the exactly converted INT8 V rows, in FP16 `mma` with FP32 accumulation. They therefore differ from +the CUDA-core kernels by the rounding of `Q'` and of the probability tile and by the FP32 summation order, so outputs are +close but not bit-identical to the previous kernels. A CTA (one query row and KV head, that is a group of up to 16 query +heads over the row's candidates) that cannot bound that rounding does not clip or round: it evaluates its heads in FP32 +on the CUDA cores (`AttendRowFp32`), which is exact. The decision is uniform across the CTA and the FP32 path starts +from the inputs, so the result of a CTA is either the tensor-core result within the bounds below or the FP32 result. + +| A CTA takes the FP32 path when | Bound that it enforces | RC model (Qwen3.8-Flash-Next DSA layers) | +|---|---|---| +| An element of `Q * key_scale` is not finite or its magnitude is above 65504 | `Q'` fits FP16 | Does not happen | +| For one of its query heads, `eps = scale * 128 * sum_d abs(Q'_d - fp16(Q'_d))` is above `kInt8TcQueryErrorBudget` (0.125) | Rounding `Q'` moves any logit of that head by at most `eps` nats (`abs(K_int8) <= 128`), so a probability by at most a factor `exp(eps)`. A query whose large terms cancel to a small logit exceeds the budget by orders of magnitude, an ordinary query does not | `eps` is 0.015 on average and 0.022 at most per head (table below) | +| A `value_scale` is not finite | The FP16 accumulator result is meaningful | Does not happen | +| Some channel has `abs(value_scale) > 2^16` and a valid candidate of one head is more than 20 nats (`kInt8TcTailDepth`) below the running maximum | A probability below `2^-29` is a subnormal of the `2^15 * p` tile and below `2^-40` it is zero. For `abs(value_scale) <= 2^16` the lost contribution is below `2^-17` of one value code; above it the check decides per CTA instead of sending every CTA to the FP32 path | The largest `value_scale` is 0.285: the check is never armed | + +The scores and the running maximum are kept in the natural-log domain and only the difference to the new maximum is +scaled by `log2(e)` for `ex2`, so the difference of two large, nearby logits is not rounded away. Prefill keeps a single +FP16 probability tile holding `2^15 * p`, as the BF16-dequantizing kernel does with a coarser BF16 tile; decode adds an +FP16 residual tile, because the signed sum of large value codes can cancel to a value far below the `2^-11` rounding of +one FP16 probability. + +The bound `eps` is proportional to the magnitude of the query. Monte Carlo over the 12 DSA layers of the RC model +(their `q_norm` weights and key scales, 4000 rows per layer and KV head, scale `1/16`, unit-RMS Gaussian queries with +partial rotary embedding rounded to BF16, then multiplied by `m`; `m = 1` is what the model produces): + +| `m` | mean `eps` of a head | largest `eps` of a head | CTAs over the budget (12 heads per CTA) | +|---|---|---|---| +| 1 | 0.015 | 0.022 | 0 | +| 2 | 0.031 | 0.045 | 0 | +| 4 | 0.061 | 0.088 | 0 | +| 8 | 0.122 | 0.179 | 69% | +| 16 | 0.245 | 0.356 | 100% | + +The budget is therefore about 8 times the magnitude the model produces. Heavy-tailed queries (Student-t with 4 degrees +of freedom) reach 48% of the CTAs at `m = 8`, queries whose channel variances differ by a log-normal factor (sigma 1) +about 1%. Workloads whose queries are 8 times larger than those of the RC model, or whose cache has +`abs(value_scale) > 2^16` and wide logit spreads, should leave the switches at `0`: the FP32 path is slower than the +previous kernels, as the measurements below show. It stays accurate: with queries 8 times larger the relative RMS error +against FP64 is 1.7e-3 (the previous kernels: 8.5e-3 for prefill, 1.7e-3 for decode). + +Cost of the FP32 path on an H200 (`benchmark_dynamic_sparse_attention.py`; 24 query / 2 KV heads, head size 256, +2048 selected candidates, context 8192; attention-kernel time, previous kernel: BF16-dequantizing prefill / CUDA-core +decode): + +| Data | Rows | Previous kernel | Tensor-core kernel | Tensor core vs previous | +|---|---|---|---|---| +| Queries like the RC model | 128 | 0.173 ms | 0.063 ms | 2.7x faster | +| | 8192 | 7.82 ms | 2.83 ms | 2.8x faster | +| Queries 8 times larger (`--query-std 11.2`, every CTA over the budget) | 128 | 0.172 ms | 3.68 ms | 21x slower | +| | 8192 | 7.82 ms | 135.6 ms | 17x slower | +| | 8 (decode) | 73.6 us | 240 us | 3.3x slower | +| `abs(value_scale)` 24k to 360k (`--value-scale-multiplier 1.2e7`), logit standard deviation about 2.4 nats: 10 to 12% of the CTAs have a candidate deeper than 20 nats | 128 | 0.172 ms | 3.61 ms | 21x slower | +| | 512 | 0.577 ms | 3.90 ms | 6.8x slower | +| | 2048 | 2.288 ms | 6.17 ms | 2.7x slower | +| The same value scales, logit standard deviation about 1.3 nats (the cache holds only rows appended by the call): no CTA is deeper than 20 nats | 8192 | 7.78 ms | 2.86 ms | 2.7x faster | +| The same value scales | 8 (decode) | 73.6 us | 13.3 us | 5.6x faster | + +A CTA on the FP32 path needs about 3.6 ms for 2048 candidates, so the kernel time is set by the slowest CTA of a launch +with one wave of CTAs and grows with the fraction of such CTAs for larger launches. A decode CTA covers fewer candidates +(64 to 128 at the automatic split on an H200 for 2048 candidates and 1 to 8 rows), so a candidate 20 nats below the +running maximum is rarer and the decode row of the last group shows no slowdown. With +`ORT_DSA_DECODE_INT8_TC_TILES_PER_SPLIT=16` (1024 candidates per CTA) the same data takes 1.8 ms instead of 13 us, with +the same RMS error (1.7e-3). The previous design of this kernel sent every CTA to the FP32 path as soon as one channel +had `abs(value_scale) > 2^16` (8192 rows: 2.76 ms became 136 ms), which the check above replaces; it still applies to +queries above the budget, as the second group of rows shows. The unit tests exercise every condition of the table +(`Int8TensorCoreQueryFoldingRoundingError_CUDA`, `Int8TensorCoreDecodeSignedValueCancellation_CUDA`, +`Int8TensorCoreLargeCommonLogit_CUDA`, `Int8TensorCoreTinyValueScale_CUDA`, +`Int8TensorCoreLargeValueScaleRandom_CUDA`, `Int8TensorCoreRandomQueryFolding_CUDA`, +`Int8PrefillQueryBeyondHalfRange_CUDA`, `Int8DecodeQueryBeyondHalfRange_CUDA`) against an FP64 host reference. + +### Rollout + +`ORT_DSA_PREFILL_INT8_TC` and `ORT_DSA_DECODE_INT8_TC` are opt-in and independent, so prefill and decode can be enabled +one at a time. The end-to-end accuracy measured for an earlier revision of these kernels (MMLU-Pro, 2400 items, +pooled difference +0.08 percentage points, 95% confidence interval -0.60 to +0.77) predates the numerics described above +and does not cover the final revision. The end-to-end timings in [Benchmark](#benchmark) are preliminary for the same +reason. Making a switch the default needs paired accuracy runs of the final revision with that switch toggled alone, +including long contexts, and a stated bound on the degradation. Packages that qualify the kernels enable them in their +session options (`ep.cuda.dsa_prefill_int8_tc` and `ep.cuda.dsa_decode_int8_tc` set to `1`). + ### Buffer alignment The wide-load fast paths assume more alignment than a tensor's element type guarantees, and tensors @@ -94,8 +202,8 @@ buffer is misaligned: | INT8 grouped decode vector staging (`ORT_DSA_INT8_VECTOR_STAGING=1`) | `key_scale`, `value_scale`, `auxiliary_key`, `auxiliary_value`: 16 bytes; `present_key`, `present_value`: 8 bytes | Scalar per-channel staging (`kVectorizedStaging=false`). | | Floating-point grouped decode and grouped prefill | Each staged K/V row: 16 bytes (checked per row in the kernel) | Element-wise copy into shared memory. | | Tensor-core BF16 prefill (`ORT_DSA_PREFILL_KERNEL=1`) | `present_key`, `present_value`: 16 bytes (8 for INT8); `key_scale`, `value_scale`: 16 bytes; `output`: 4 bytes | Generic fused/split or grouped prefill kernels. | -| Tensor-core INT8-cache prefill (`ORT_DSA_PREFILL_INT8_TC=1`) | In addition to the row above: `present_key`, `present_value`: 16 bytes; `output`: 8 bytes | `DynamicSparseAttentionPrefillKernel`. | -| Tensor-core INT8-cache decode (`ORT_DSA_DECODE_INT8_TC=1`) | `prepared_query`, `present_key`, `present_value`, `key_scale`, `value_scale`, attention workspace: 16 bytes | The CUDA-core grouped decode kernel. | +| Tensor-core INT8-cache prefill (opt-in, `ORT_DSA_PREFILL_INT8_TC=1`) | In addition to the row above: `present_key`, `present_value`: 16 bytes; `output`: 8 bytes | `DynamicSparseAttentionPrefillKernel`. | +| Tensor-core INT8-cache decode (opt-in, `ORT_DSA_DECODE_INT8_TC=1`) | `prepared_query`, `present_key`, `present_value`, `key_scale`, `value_scale`, attention workspace: 16 bytes | The CUDA-core grouped decode kernel. | Buffers allocated by ONNX Runtime are 256-byte aligned, so only IOBinding with offset buffers takes a fallback. The `*MisalignedBuffers_CUDA` tests cover each case. @@ -111,28 +219,70 @@ only as fallbacks: | `ORT_DSA_INT8_VECTOR_STAGING=0` | about +0.3 ms (+2%) slower | | `ORT_DSA_SHARED_SCALE_MERGE=0` | about +0.05 to +0.1 ms (+0.3 to 0.5%) slower | -Prefill with the INT8 KV cache (24 query / 2 KV heads, head size 256, 12 full-attention layers), H200: -`DynamicSparseAttentionInt8TcKernel` takes 2.15 ms per launch instead of 6.1 ms in a standalone -harness (8192 rows, 2048 candidates, 2.83x) and 2.8 ms instead of 7.9 ms per 8192-row launch in the -model. The previous kernel spent 2.8% of its instructions on `mma` and the rest on in-kernel -dequantization and index handling. Time to first token in the Qwen3.8-Flash-Next packages (one round) -drops by 4.0 to 4.8% at 2048, 8.4 to 9.5% at 8192 and 7.6 to 9.4% at 32768 prompt tokens; decode is -unchanged. The RMS error against an FP64 reference is 1.6e-3 (previous kernel 1.7e-3). -`ORT_DSA_PREFILL_INT8_TC=0` restores the previous kernel. The test -`DynamicSparseAttentionTest.Int8CachePrefillTensorCoreKernel_CUDA` compares both kernels with a -reference over head counts, candidate counts and invalid indices. - -Decode with the INT8 KV cache (`sequence_length <= 8`, 12 query heads per KV head, head size 256), H200: the -CUDA-core grouped decode kernel was instruction-issue and shared-load bound on its per-candidate serial chain, -70.3 us per layer plus 5.3 us for the merge (8 rows, 2048 selected candidates, flat from 2K to 32K context). The -split-K tensor-core kernel takes 9.5 us plus 3.5 us for the merge. In the Qwen3.8-Flash-Next packages the -speculative-decoding cycle (S = 8) is 4.6 to 6.8% shorter at 2K, 8K and 32K context. The tensor-core kernel -accumulates Q.K and P.V in a different FP32 order, so outputs are close but not bit-identical to the previous kernel; -`ORT_DSA_DECODE_INT8_TC=0` restores it exactly. The tests `Int8CacheDecodeTensorCoreKernel_CUDA` (reference -comparison over splits, candidate counts and invalid indices), `Int8DecodeQueryBeyondHalfRange_CUDA` (FP32 fallback -of the splits) and `DecodeSessionOptions_CUDA` (switch, tuning knob and invalid values) cover it. +The INT8 tensor-core kernels (opt-in, see [Rollout](#rollout)) are measured at the operator level with +[benchmark_dynamic_sparse_attention.py](../../../onnxruntime/test/python/transformers/benchmark_dynamic_sparse_attention.py) +against the kernels that they replace (the BF16-dequantizing prefill and the CUDA-core grouped decode). The script builds +a single-node `DynamicSparseAttention` model per case (24 query / 2 KV heads, head size 256, INT8 cache, scale 1/16, +2048 selected candidates per row: the shapes of the 12 full-attention layers of Qwen3.8-Flash-Next) and one session per +kernel configuration, because the `ep.cuda.dsa_*` options are read when the kernel is created. The time is the GPU time +of the kernels that one `Run` launches, taken from the CUPTI-based torch profiler, so host launch overhead does not +enter: "attention" sums the `DynamicSparseAttention` kernels including the decode merge kernel, "node" all GPU activity +of the node (cache append, query preparation, ...). Each round measures every configuration once, the order alternates +between rounds, the tables show the median of 7 rounds, and the speedup is the median of the paired ratios of the rounds. +The error is the RMS error of the BF16 output relative to the RMS of an FP64 reference over 64 evenly spaced rows. The +data are synthetic: BF16 queries with standard deviation 1.4 (the queries of the RC model after `q_norm` have a per-head +RMS of 1.3 to 1.6), INT8 cache codes with standard deviation 40, and key and value scales drawn from the ranges of that +model (`--rc-model` reads the scales of a model's node instead). H200, CUDA 13.0, the automatic tiles per split: + +```bash +python benchmark_dynamic_sparse_attention.py --rows 1 8 --context 2048 8192 32768 --tiles-per-split 0 1 2 4 8 16 +python benchmark_dynamic_sparse_attention.py --rows 128 2048 8192 --context 8192 32768 +``` + +Decode (`sequence_length <= 8`): + +| Rows | Context | Attention, previous | Attention, tensor core | Speedup | Node, previous | Node, tensor core | RMS error, previous / tensor core | +|---|---|---|---|---|---|---|---| +| 1 | 2048 | 31.1 us | 9.5 us | 3.30x | 34.1 us | 12.5 us | 1.59e-3 / 1.62e-3 | +| 1 | 8192 | 31.2 us | 9.5 us | 3.30x | 34.2 us | 12.5 us | 1.55e-3 / 1.57e-3 | +| 1 | 32768 | 31.2 us | 9.5 us | 3.30x | 34.1 us | 12.5 us | 1.67e-3 / 1.72e-3 | +| 8 | 2048 | 73.6 us | 13.2 us | 5.58x | 76.7 us | 16.5 us | 1.64e-3 / 1.69e-3 | +| 8 | 8192 | 73.6 us | 13.3 us | 5.57x | 76.8 us | 16.6 us | 1.62e-3 / 1.66e-3 | +| 8 | 32768 | 73.7 us | 13.2 us | 5.57x | 77.0 us | 16.6 us | 1.68e-3 / 1.73e-3 | + +Prefill: + +| Rows | Context | Attention, previous | Attention, tensor core | Speedup | Node, previous | Node, tensor core | RMS error, previous / tensor core | +|---|---|---|---|---|---|---|---| +| 128 | 8192 | 0.173 ms | 0.063 ms | 2.72x | 0.179 ms | 0.070 ms | 4.19e-3 / 1.71e-3 | +| 128 | 32768 | 0.191 ms | 0.064 ms | 2.99x | 0.199 ms | 0.072 ms | 4.25e-3 / 1.70e-3 | +| 2048 | 8192 | 2.287 ms | 0.832 ms | 2.75x | 2.350 ms | 0.894 ms | 4.25e-3 / 1.71e-3 | +| 2048 | 32768 | 2.294 ms | 0.832 ms | 2.75x | 2.356 ms | 0.895 ms | 4.21e-3 / 1.70e-3 | +| 8192 | 8192 | 7.818 ms | 2.830 ms | 2.76x | 8.057 ms | 3.069 ms | 2.69e-3 / 1.66e-3 | +| 8192 | 32768 | 8.877 ms | 3.200 ms | 2.77x | 9.116 ms | 3.439 ms | 4.29e-3 / 1.71e-3 | + +The previous prefill kernel spent 2.8% of its instructions on `mma` and the rest on in-kernel dequantization and index +handling, and the previous grouped decode kernel was instruction-issue and shared-load bound on its per-candidate +serial chain. The decode time does not depend on the context length (the selected candidates are gathered). The error +includes the rounding of the BF16 output, whose floor is about 1.5e-3 relative to the RMS: both decode kernels and the +tensor-core prefill are near it, and the BF16-dequantizing prefill, which rounds the dequantized keys and values and the +probabilities to BF16, has 1.6 to 2.5 times the error of the tensor-core prefill. The tensor-core kernels accumulate Q.K +and P.V in a different FP32 order, so their outputs are close but not bit-identical to the previous kernels. + +End-to-end numbers are preliminary. In the Qwen3.8-Flash-Next packages, an earlier revision of these kernels +(one round of consecutive A/B runs, not alternated) shortened the time to first token by 4.0 to 4.8% at 2048, 8.4 to 9.5% +at 8192 and 7.6 to 9.4% at 32768 prompt tokens, the speculative-decoding cycle (S = 8) by 4.6 to 6.8% at 2K, 8K and 32K +context, and the 8192-row prefill launch from 7.9 to 2.8 ms. They have not been reproduced for the current revision of +the kernels, whose numerics changed (see [Rollout](#rollout)). + +`ORT_DSA_PREFILL_INT8_TC=0` and `ORT_DSA_DECODE_INT8_TC=0` (the defaults) run the previous kernels exactly. The tests +`Int8CachePrefillTensorCoreKernel_CUDA` and `Int8CacheDecodeTensorCoreKernel_CUDA` compare the tensor-core kernels with a +reference over head counts, candidate counts, splits and invalid indices; `Int8DecodeQueryBeyondHalfRange_CUDA` covers +the FP32 fallback of the splits, `Int8TensorCoreDecodeTilesPerSplitCaps_CUDA` the split caps, and +`PrefillSessionOptions_CUDA` and `DecodeSessionOptions_CUDA` the switches, the tuning knob and invalid values. `Int8TensorCoreLargeSoftmaxScale_CUDA` (the `scale > 1` fallback) and -`Int8TensorCoreTinyProbabilityLargeValueScale_CUDA` (the FP16 probability range) cover prefill and decode. +`Int8TensorCoreTinyProbabilityLargeValueScale_CUDA` (the FP16 probability range) cover prefill and decode, and the tests +listed under [Numerics and fallbacks](#numerics-and-fallbacks) the guards. The corresponding indexer switch results are in [../packed_sparse_attention_indexer.md](../packed_sparse_attention_indexer.md#benchmark). diff --git a/onnxruntime/test/python/transformers/benchmark_dynamic_sparse_attention.py b/onnxruntime/test/python/transformers/benchmark_dynamic_sparse_attention.py new file mode 100644 index 0000000000000..28942551e5be9 --- /dev/null +++ b/onnxruntime/test/python/transformers/benchmark_dynamic_sparse_attention.py @@ -0,0 +1,422 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +""" +Op-level benchmark of the INT8 KV-cache DynamicSparseAttention CUDA kernels. + +The shapes are those of the 12 full-attention layers of Qwen3.8-Flash-Next: 24 query heads, 2 KV heads, head size 256, +INT8 KV cache with per-channel scales, softmax scale 1/16 and 2048 selected candidates per query row. Each +(rows, context length) case becomes a single-node model, and each kernel configuration becomes its own session, +because the ep.cuda.dsa_* options are read when the kernel is created: + + rows >= 9 (prefill) legacy: ep.cuda.dsa_prefill_int8_tc=0, the BF16-dequantizing kernel + tc: ep.cuda.dsa_prefill_int8_tc=1, the tensor-core kernel + rows <= 8 (decode) legacy: ep.cuda.dsa_decode_int8_tc=0, the CUDA-core grouped kernel + tc: ep.cuda.dsa_decode_int8_tc=1, the split-K tensor-core kernel, once for each + --tiles-per-split value (0 = automatic) + +The time of a configuration is the GPU time of the CUDA kernels one Run launches, measured with the CUPTI-based torch +profiler, so host launch overhead does not enter. "attention" sums the DynamicSparseAttention kernels (including the +decode merge kernel), "total" all GPU activity of the node (cache append, query preparation, ...). A round measures +every configuration once, back to back, and the order alternates between rounds. The speedup of a configuration is the +median over rounds of the paired ratio to the first (legacy) configuration, which cancels clock and temperature drift. +The "rms err" column is the RMS error of the BF16 output relative to the RMS of an FP64 reference, computed over +--reference-rows evenly spaced query rows (0 disables it). It is the error against the exact result, not the +difference to the legacy kernel, which has errors of its own. + +Example (a CUDA plugin EP build needs ORT_TEST_CUDA_PLUGIN_EP=1 and, unless the wheel carries the freshly built +library, ORT_CUDA_PLUGIN_PATH): + + python benchmark_dynamic_sparse_attention.py --rows 8 --context 2048 8192 32768 --tiles-per-split 0 1 2 4 8 16 + python benchmark_dynamic_sparse_attention.py --rows 128 2048 8192 --context 8192 32768 --json prefill.json + +The key and value scales are synthetic (uniform in the ranges of the Qwen3.8-Flash-Next layers) unless --rc-model names +a model whose DynamicSparseAttention node supplies them. The multipliers below reproduce the slow FP32 fallback of the +tensor-core kernels (queries or value scales far outside the usual range). +""" + +import argparse +import json +import platform +import statistics +import sys + +import numpy as np +import onnx +import torch +from cuda_plugin_ep_helper import CUDA_PLUGIN_EP_NAME, ensure_cuda_plugin_ep_registered +from onnx import TensorProto, helper, numpy_helper + +import onnxruntime as ort + +NUM_HEADS = 24 +KV_HEADS = 2 +HEAD_SIZE = 256 +SOFTMAX_SCALE = 0.0625 +MAX_DECODE_ROWS = 8 +TORCH_TO_ONNX = { + torch.bfloat16: TensorProto.BFLOAT16, + torch.int8: TensorProto.INT8, + torch.int32: TensorProto.INT32, +} + + +def load_rc_scales(model_path, layer): + """Returns the INT8 key and value scales of the layer-th DynamicSparseAttention node of a model.""" + model = onnx.load(model_path, load_external_data=False) + initializers = {initializer.name: initializer for initializer in model.graph.initializer} + nodes = [node for node in model.graph.node if node.op_type == "DynamicSparseAttention"] + if layer >= len(nodes): + raise ValueError(f"{model_path} has {len(nodes)} DynamicSparseAttention nodes, --layer {layer} is out of range") + return tuple( + numpy_helper.to_array(initializers[nodes[layer].input[index]]).astype(np.float32) for index in (18, 19) + ) + + +def make_scales(args): + if args.rc_model: + key_scale, value_scale = load_rc_scales(args.rc_model, args.layer) + else: + rng = np.random.default_rng(args.seed) + key_scale = rng.uniform(0.002, 0.07, KV_HEADS * HEAD_SIZE).astype(np.float32) + value_scale = rng.uniform(0.002, 0.03, KV_HEADS * HEAD_SIZE).astype(np.float32) + return key_scale * np.float32(args.key_scale_multiplier), value_scale * np.float32(args.value_scale_multiplier) + + +def make_model(rows, context, max_selected, key_scale, value_scale): + def value_info(name, element_type, shape): + return helper.make_tensor_value_info(name, element_type, shape) + + cache_shape = [1, KV_HEADS, context, HEAD_SIZE] + graph = helper.make_graph( + [ + helper.make_node( + "DynamicSparseAttention", + [ + "query", + "key", + "value", + "past_key", + "past_value", + "", + "", + "selected_indices", + "selected_counts", + "seqlens_k", + "total_sequence_length", + *[""] * 7, + "key_scale", + "value_scale", + ], + ["output", "present_key", "present_value"], + domain="com.microsoft", + num_heads=NUM_HEADS, + kv_num_heads=KV_HEADS, + scale=SOFTMAX_SCALE, + ) + ], + "dynamic_sparse_attention", + [ + value_info("query", TensorProto.BFLOAT16, [1, rows, NUM_HEADS * HEAD_SIZE]), + value_info("key", TensorProto.BFLOAT16, [1, rows, KV_HEADS * HEAD_SIZE]), + value_info("value", TensorProto.BFLOAT16, [1, rows, KV_HEADS * HEAD_SIZE]), + value_info("past_key", TensorProto.INT8, cache_shape), + value_info("past_value", TensorProto.INT8, cache_shape), + value_info("selected_indices", TensorProto.INT32, [rows, max_selected]), + value_info("selected_counts", TensorProto.INT32, [rows]), + value_info("seqlens_k", TensorProto.INT32, [1]), + ], + [ + value_info("output", TensorProto.BFLOAT16, [1, rows, NUM_HEADS * HEAD_SIZE]), + value_info("present_key", TensorProto.INT8, cache_shape), + value_info("present_value", TensorProto.INT8, cache_shape), + ], + [ + numpy_helper.from_array(np.array(context, dtype=np.int32), "total_sequence_length"), + numpy_helper.from_array(key_scale, "key_scale"), + numpy_helper.from_array(value_scale, "value_scale"), + ], + ) + model = helper.make_model( + graph, opset_imports=[helper.make_opsetid("", 21), helper.make_opsetid("com.microsoft", 1)] + ) + model.ir_version = 10 + return model.SerializeToString() + + +def make_tensors(rows, context, max_selected, args): + """The query rows are the last `rows` positions of a context of `context` tokens; each selects random positions.""" + generator = torch.Generator(device="cuda") + generator.manual_seed(args.seed) + + def normal(*shape, std): + return torch.randn(*shape, device="cuda", generator=generator) * std + + def cache_codes(): + return normal(1, KV_HEADS, context, HEAD_SIZE, std=args.cache_code_std).clamp(-127, 127).round().to(torch.int8) + + positions = torch.arange(context - rows, context, device="cuda") + counts = (positions + 1).clamp(max=max_selected).to(torch.int32) + selected = torch.empty(rows, max_selected, dtype=torch.int32, device="cuda") + columns = torch.arange(context, device="cuda") + for start in range(0, rows, 1024): + stop = min(start + 1024, rows) + noise = torch.rand(stop - start, context, device="cuda", generator=generator) + noise[columns[None, :] > positions[start:stop, None]] = -1.0 + selected[start:stop] = noise.topk(max_selected, dim=1).indices.to(torch.int32) + selected[torch.arange(max_selected, device="cuda")[None, :] >= counts[:, None]] = -1 + return { + "query": normal(1, rows, NUM_HEADS * HEAD_SIZE, std=args.query_std).to(torch.bfloat16), + "key": normal(1, rows, KV_HEADS * HEAD_SIZE, std=1.0).to(torch.bfloat16), + "value": normal(1, rows, KV_HEADS * HEAD_SIZE, std=1.0).to(torch.bfloat16), + "past_key": cache_codes(), + "past_value": cache_codes(), + "selected_indices": selected, + "selected_counts": counts, + "seqlens_k": torch.tensor([context - 1], dtype=torch.int32, device="cuda"), + } + + +def select_plugin_device(device_id): + """The device of the CUDA plugin EP when ORT_TEST_CUDA_PLUGIN_EP=1 selects it, otherwise None.""" + if not ensure_cuda_plugin_ep_registered(): + return None + for device in ort.get_ep_devices(): + if device.ep_name == CUDA_PLUGIN_EP_NAME: + found = device.ep_options.get("device_id", device.ep_metadata.get("cuda_device_id")) + if found is None or int(found) == device_id: + return device + return None + + +def create_session(model_bytes, options, args): + session_options = ort.SessionOptions() + for key, value in options.items(): + session_options.add_session_config_entry(key, value) + plugin_device = select_plugin_device(args.device) + if plugin_device is not None: + session_options.add_provider_for_devices([plugin_device], {}) + return ort.InferenceSession(model_bytes, session_options) + providers = [("CUDAExecutionProvider", {"device_id": args.device})] + return ort.InferenceSession(model_bytes, session_options, providers=providers) + + +def create_binding(session, tensors, output, args): + binding = session.io_binding() + for name, tensor in tensors.items(): + binding.bind_input( + name, "cuda", args.device, TORCH_TO_ONNX[tensor.dtype], list(tensor.shape), tensor.data_ptr() + ) + # The cache is updated in place, like the model does with its past_key_values buffers. + binding.bind_output("output", "cuda", args.device, TensorProto.BFLOAT16, list(output.shape), output.data_ptr()) + for name, cache in (("present_key", tensors["past_key"]), ("present_value", tensors["past_value"])): + binding.bind_output(name, "cuda", args.device, TensorProto.INT8, list(cache.shape), cache.data_ptr()) + return binding + + +def reference_attention(tensors, key_scale, value_scale, rows_to_check): + """FP64 attention of the given query rows over the cache the operator updated in place (new rows appended).""" + scales = [ + torch.tensor(scale, device="cuda", dtype=torch.float64).view(KV_HEADS, 1, HEAD_SIZE) + for scale in (key_scale, value_scale) + ] + keys = tensors["past_key"][0].double() * scales[0] + values = tensors["past_value"][0].double() * scales[1] + query = tensors["query"][0].double().view(-1, NUM_HEADS, HEAD_SIZE) + group = NUM_HEADS // KV_HEADS + reference = torch.empty(len(rows_to_check), NUM_HEADS, HEAD_SIZE, dtype=torch.float64, device="cuda") + for position, row in enumerate(rows_to_check): + selected = tensors["selected_indices"][row, : int(tensors["selected_counts"][row])].long() + for kv_head in range(KV_HEADS): + heads = slice(kv_head * group, (kv_head + 1) * group) + scores = query[row, heads] @ keys[kv_head, selected].T * SOFTMAX_SCALE + reference[position, heads] = torch.softmax(scores, dim=-1) @ values[kv_head, selected] + return reference + + +def relative_rms_error(result, reference): + return float((result - reference).pow(2).mean().sqrt() / reference.pow(2).mean().sqrt()) + + +def short_kernel_name(name): + head = name.replace("(anonymous namespace)::", "").split("(")[0].removeprefix("void ").strip() + base, separator, template_arguments = head.partition("<") + return base.split("::")[-1] + separator + template_arguments + + +def run_and_profile(session, binding, iterations): + """Returns the GPU time per Run in microseconds: {kernel name: total}, summed over the profiled runs.""" + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as profile: + for _ in range(iterations): + session.run_with_iobinding(binding) + binding.synchronize_outputs() + torch.cuda.synchronize() + times = {} + for event in profile.events(): + if event.device_type == torch.autograd.DeviceType.CUDA: + name = short_kernel_name(event.name) + times[name] = times.get(name, 0.0) + event.device_time / iterations + return times + + +def configurations(rows, args): + if rows <= MAX_DECODE_ROWS: + configs = [("legacy (CUDA-core)", {"ep.cuda.dsa_decode_int8_tc": "0"})] + for tiles in args.tiles_per_split: + options = { + "ep.cuda.dsa_decode_int8_tc": "1", + "ep.cuda.dsa_decode_int8_tc_tiles_per_split": str(tiles), + } + configs.append((f"tc tiles_per_split={'auto' if tiles == 0 else tiles}", options)) + return configs + return [ + ("legacy (BF16 dequantizing)", {"ep.cuda.dsa_prefill_int8_tc": "0"}), + ("tc", {"ep.cuda.dsa_prefill_int8_tc": "1"}), + ] + + +def format_range(values, digits): + return f"{statistics.median(values):.{digits}f} [{min(values):.{digits}f}, {max(values):.{digits}f}]" + + +def run_case(rows, context, args, key_scale, value_scale): + max_selected = min(args.selected, context) + model_bytes = make_model(rows, context, max_selected, key_scale, value_scale) + tensors = make_tensors(rows, context, max_selected, args) + output = torch.empty(1, rows, NUM_HEADS * HEAD_SIZE, dtype=torch.bfloat16, device="cuda") + configs = configurations(rows, args) + iterations = args.iterations or (args.decode_iterations if rows <= MAX_DECODE_ROWS else args.prefill_iterations) + + rows_to_check = sorted({int(row) for row in np.linspace(0, rows - 1, min(rows, args.reference_rows)).round()}) + sessions, bindings, reference, rms_error = [], [], None, [] + for _, options in configs: + session = create_session(model_bytes, options, args) + binding = create_binding(session, tensors, output, args) + for _ in range(args.warmup): + session.run_with_iobinding(binding) + binding.synchronize_outputs() + torch.cuda.synchronize() + if args.reference_rows > 0: + if reference is None: + # The first run has appended the new key/value rows to the cache; later runs rewrite the same values. + reference = reference_attention(tensors, key_scale, value_scale, rows_to_check).view( + len(rows_to_check), -1 + ) + rms_error.append(relative_rms_error(output[0, rows_to_check].float(), reference.float())) + else: + rms_error.append(float("nan")) + sessions.append(session) + bindings.append(binding) + + attention = [[] for _ in configs] + total = [[] for _ in configs] + kernels = [{} for _ in configs] + for round_index in range(args.rounds): + order = range(len(configs)) if round_index % 2 == 0 else reversed(range(len(configs))) + for index in order: + for _ in range(args.warmup): + sessions[index].run_with_iobinding(bindings[index]) + times = run_and_profile(sessions[index], bindings[index], iterations) + attention[index].append(sum(t for name, t in times.items() if "DynamicSparseAttention" in name)) + total[index].append(sum(times.values())) + kernels[index] = times + + kind = "decode" if rows <= MAX_DECODE_ROWS else "prefill" + unit, scale, digits = ("us", 1.0, 2) if kind == "decode" else ("ms", 1e-3, 3) + print( + f"\n{kind}: rows={rows}, context={context}, selected={max_selected}, {iterations} runs x {args.rounds} rounds" + ) + print(f"| configuration | attention {unit} | total {unit} | speedup | rms err |") + print("|---|---|---|---|---|") + results = [] + for index, (label, options) in enumerate(configs): + ratios = [base / value for base, value in zip(attention[0], attention[index], strict=True)] + attention_scaled = [value * scale for value in attention[index]] + total_scaled = [value * scale for value in total[index]] + print( + f"| {label} | {format_range(attention_scaled, digits)} | {format_range(total_scaled, digits)} " + f"| {format_range(ratios, 2)}x | {rms_error[index]:.2e} |" + ) + results.append( + { + "label": label, + "options": options, + "attention_us": attention[index], + "total_us": total[index], + "kernels_us": kernels[index], + "rms_error": rms_error[index], + } + ) + if args.breakdown: + for name, value in sorted(kernels[index].items(), key=lambda item: -item[1])[:5]: + print(f"| {name[:90]} | {value:.2f} us | | | |") + return { + "kind": kind, + "rows": rows, + "context": context, + "selected": max_selected, + "iterations": iterations, + "configurations": results, + } + + +def environment(): + properties = torch.cuda.get_device_properties(torch.cuda.current_device()) + return { + "onnxruntime": ort.__version__, + "onnxruntime_build_info": ort.get_build_info(), + "torch": torch.__version__, + "gpu": properties.name, + "sm_count": properties.multi_processor_count, + "compute_capability": f"{properties.major}.{properties.minor}", + "python": platform.python_version(), + "command": " ".join(sys.argv), + } + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--rows", type=int, nargs="+", default=[8, 128, 8192], help="query rows; <= 8 selects decode") + parser.add_argument("--context", type=int, nargs="+", default=[8192], help="total sequence lengths") + parser.add_argument("--selected", type=int, default=2048, help="selected candidates per row") + parser.add_argument( + "--tiles-per-split", type=int, nargs="+", default=[0], help="decode tc tiles per split (0=auto)" + ) + parser.add_argument("--rounds", type=int, default=7, help="alternating rounds over all configurations") + parser.add_argument("--iterations", type=int, default=None, help="profiled runs per round (default by kind)") + parser.add_argument("--prefill-iterations", type=int, default=5) + parser.add_argument("--decode-iterations", type=int, default=200) + parser.add_argument("--warmup", type=int, default=3, help="unprofiled runs before each measurement") + parser.add_argument("--device", type=int, default=0) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--reference-rows", type=int, default=64, help="query rows checked against FP64 (0=skip)") + parser.add_argument("--rc-model", default="", help="model whose DynamicSparseAttention node supplies the scales") + parser.add_argument("--layer", type=int, default=0, help="index of that DynamicSparseAttention node") + parser.add_argument("--query-std", type=float, default=1.4, help="BF16 query values ~ N(0, std^2); 1.4 is typical") + parser.add_argument("--cache-code-std", type=float, default=40.0, help="INT8 cache codes ~ N(0, std^2), clipped") + parser.add_argument("--key-scale-multiplier", type=float, default=1.0) + parser.add_argument("--value-scale-multiplier", type=float, default=1.0, help="1.2e7 gives scales above 2^16") + parser.add_argument("--breakdown", action="store_true", help="also print the largest kernels of each configuration") + parser.add_argument("--json", default="", help="write the raw per-round timings and the environment here") + return parser.parse_args() + + +def main(): + args = parse_args() + torch.cuda.set_device(args.device) + key_scale, value_scale = make_scales(args) + info = environment() + print(json.dumps(info, indent=2)) + cases = [] + for rows in args.rows: + for context in args.context: + if rows > context: + continue + cases.append(run_case(rows, context, args, key_scale, value_scale)) + if args.json: + with open(args.json, "w") as file: + json.dump({"environment": info, "arguments": vars(args), "cases": cases}, file, indent=1) + + +if __name__ == "__main__": + main() From e8b3ad2232663d1873acccd6944b88476e74b4db Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 31 Oct 2026 20:59:30 +0000 Subject: [PATCH 13/14] [CUDA] HyperConnectionLowRankMix docs: label the end-to-end timings as preliminary The early PDL trigger and branch-free sigmoid paragraphs stated their end-to-end timings as settled facts. The 1.4-1.7% decode-cycle and 0.3-1.1% TTFT figures come from a single A/B on an earlier revision, and the TTFT change is about the size of the repeat-to-repeat noise, so say that and point to the kernel times as the firmer evidence. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cuda/hyper_connection_low_rank_mix.md | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/docs/contrib_ops/cuda/hyper_connection_low_rank_mix.md b/docs/contrib_ops/cuda/hyper_connection_low_rank_mix.md index 8c0ec5331d088..6d772a36c8397 100644 --- a/docs/contrib_ops/cuda/hyper_connection_low_rank_mix.md +++ b/docs/contrib_ops/cuda/hyper_connection_low_rank_mix.md @@ -45,7 +45,8 @@ is launched with PDL, its prologue (constant-weight prefetch) then overlaps with the gate-mul kernel, while its reads of the gate-mul output still wait for `griddepcontrol.wait`. The trigger has no effect for plain launches, changes no result and has no switch. Together with the same early trigger in the fpA_intB GEMV -kernels (MatMulNBits) it shortens a Qwen3.8-Flash-Next `S = 8` decode cycle by 1.4-1.7%. +kernels (MatMulNBits) it shortened a Qwen3.8-Flash-Next `S = 8` decode cycle by 1.4-1.7% in a +preliminary end-to-end A/B on an H200 (one round, not alternated, no confidence interval). `SigmoidF32` evaluates `expf(-|v|)` once and selects the numerator, `e = expf(-|v|); (v >= 0 ? 1 : e) / (1 + e)`. The textbook two-branch form passes the @@ -54,9 +55,12 @@ bit-identical for all 2^32 `float` values (checked exhaustively in a standalone no fast-math). The branch itself was the cost: in the prefill pre-mix neighbouring elements have different signs, so nearly every warp executed both sides. For hidden 2560, 4 branches and an 8192-token prefill chunk on an H200, `PreMixVec8Kernel` takes -82.7 µs instead of 112.1 µs per launch and `ScaledSiluKernel` 6.1 µs instead of 7.3 µs, -which shortens the TTFT by 0.3-1.1%. The fused decode kernels call the same function for -at most eight rows, and decode time is unchanged within the noise of 0.3%. +82.7 µs instead of 112.1 µs per launch and `ScaledSiluKernel` 6.1 µs instead of 7.3 µs. +The TTFT shortened by 0.3-1.1% in a preliminary end-to-end A/B (median of 5 repetitions, +faster in 9 of 9 package and context cells), which is about the size of the roughly 1% +noise between repetitions, so the kernel times are the firmer evidence. The fused decode +kernels call the same function for at most eight rows, and decode time is unchanged within +the noise of 0.3%. The rank-320 decode test (`HyperConnectionLowRankMixTest.FusedDecodeShapeRank320`) covers 1, 2, 3, 4, 7 and 8 rows. From d60b0f4b639c4ca40ab4d0ba06a44606312fd4e1 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 31 Oct 2026 21:06:58 +0000 Subject: [PATCH 14/14] [CUDA] DynamicSparseAttention docs: state that the end-to-end timings lack acceptance measurements The end-to-end timings of the INT8 tensor-core kernels are single A/B rounds of an earlier revision. Say that they come without raw timings or confidence intervals, and that a shorter speculative-decoding cycle is not better user-visible throughput by itself because the accepted tokens per cycle and the time per accepted token were not measured. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/contrib_ops/cuda/dynamic_sparse_attention.md | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/docs/contrib_ops/cuda/dynamic_sparse_attention.md b/docs/contrib_ops/cuda/dynamic_sparse_attention.md index 337dd5e4d9a87..5635d7c9b93d5 100644 --- a/docs/contrib_ops/cuda/dynamic_sparse_attention.md +++ b/docs/contrib_ops/cuda/dynamic_sparse_attention.md @@ -272,8 +272,10 @@ and P.V in a different FP32 order, so their outputs are close but not bit-identi End-to-end numbers are preliminary. In the Qwen3.8-Flash-Next packages, an earlier revision of these kernels (one round of consecutive A/B runs, not alternated) shortened the time to first token by 4.0 to 4.8% at 2048, 8.4 to 9.5% at 8192 and 7.6 to 9.4% at 32768 prompt tokens, the speculative-decoding cycle (S = 8) by 4.6 to 6.8% at 2K, 8K and 32K -context, and the 8192-row prefill launch from 7.9 to 2.8 ms. They have not been reproduced for the current revision of -the kernels, whose numerics changed (see [Rollout](#rollout)). +context, and the 8192-row prefill launch from 7.9 to 2.8 ms. They are single A/B rounds without raw timings or +confidence intervals and have not been reproduced for the current revision of the kernels, whose numerics changed (see +[Rollout](#rollout)). A shorter speculative-decoding cycle is not better user-visible throughput by itself: the accepted +tokens per cycle and the time per accepted token were not measured. `ORT_DSA_PREFILL_INT8_TC=0` and `ORT_DSA_DECODE_INT8_TC=0` (the defaults) run the previous kernels exactly. The tests `Int8CachePrefillTensorCoreKernel_CUDA` and `Int8CacheDecodeTensorCoreKernel_CUDA` compare the tensor-core kernels with a