[CUDA] Tensor-core INT8-KV DynamicSparseAttention prefill and decode, branch-free hyper-connection sigmoid

docs/contrib_ops/cuda/hyper_connection_low_rank_mix.md

@@ -37,6 +37,31 @@
   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 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
+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.
+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.
 

include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h

@@ -510,6 +510,29 @@ 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.
+/// "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";
+
+/// "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";
+
+/// 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";
+
 /// 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

onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc

@@ -105,6 +105,15 @@ DynamicSparseAttention<T>::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_ = 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<int>::max());
   strict_validation_ = ParseEnvironmentVariableWithDefault<bool>(kStrictValidation, false);
   attention_mode_ = ParseAttentionMode(
       info.GetAttrOrDefault<std::string>("attention_mode", "selected_only"));
@@ -203,6 +212,10 @@ Status DynamicSparseAttention<T>::ComputeInternal(OpKernelContext* context) cons
 
   using CudaT = typename OrtToCudaType<T>::type;
   DynamicSparseAttentionData<CudaT> 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<const CudaT*>(query->Data<T>());
   data.key = key == nullptr ? nullptr : reinterpret_cast<const CudaT*>(key->Data<T>());
   data.value = value == nullptr ? nullptr : reinterpret_cast<const CudaT*>(value->Data<T>());

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,10 @@ 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_;
+  bool decode_int8_tc_;
+  int decode_int8_tc_tiles_per_split_;
   DynamicSparseAttentionMode attention_mode_;
   DynamicSparseAttentionKvSource selected_kv_source_;
 };

onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu

@@ -15,6 +15,7 @@
 #include <type_traits>
 
 #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"
 
@@ -1806,6 +1807,31 @@ int GetGroupedDecodeSplitCount(const DynamicSparseAttentionParameters& parameter
   return static_cast<int>(std::min({candidate_splits, occupancy_splits, max_splits}));
 }
 
+// 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 resident_ctas) {
+  const int64_t tiles = (static_cast<int64_t>(parameters.max_selected) + dsa_prefill::kTile - 1) / dsa_prefill::kTile;
+  const int64_t row_kv_heads =
+      static_cast<int64_t>(parameters.batch_size) * parameters.sequence_length * parameters.kv_num_heads;
+  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<int>(std::max<int64_t>(1, std::min<int64_t>({splits, workspace_splits, 32})));
+}
+
 size_t GetValidationHashEntries(const DynamicSparseAttentionParameters& parameters) {
   if (parameters.max_selected == 0) {
     return 0;
@@ -2073,9 +2099,10 @@ 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<T, __nv_bfloat16>) {
     if (!use_grouped_decode) {
-      const bool enabled = ParseEnvironmentVariableWithDefault<int>("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;
@@ -2101,10 +2128,84 @@ Status LaunchDynamicSparseAttention(
           parameters.num_heads % parameters.kv_num_heads == 0 &&
           parameters.num_heads / parameters.kv_num_heads <= dsa_prefill::kMaxGroup &&
           static_cast<size_t>(optin_shared) >= dsa_prefill::kSharedBytes;
+      if (use_prefill_kernel) {
+        cudaFuncAttributes attributes{};
+        const auto kernel = data.quantized_cache ? dsa_prefill::DynamicSparseAttentionPrefillKernel<true>
+                                                 : dsa_prefill::DynamicSparseAttentionPrefillKernel<false>;
+        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, 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 && parameters.scale <= dsa_prefill::kInt8TcMaxScale &&
+          IsAligned(data.present_key, 16) && IsAligned(data.present_value, 16) && IsAligned(data.output, 8) &&
+          static_cast<size_t>(optin_shared) >= dsa_prefill::kInt8SharedBytes;
+      if (use_int8_tc_prefill) {
+        cudaFuncAttributes attributes{};
+        CUDA_RETURN_IF_ERROR(cudaFuncGetAttributes(&attributes, dsa_prefill::DynamicSparseAttentionInt8TcKernel<false>));
+        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;
+  int int8_tc_decode_resident_ctas = 0;
+  if constexpr (std::is_same_v<T, __nv_bfloat16>) {
+    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.
+      // 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.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 &&
+          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<size_t>(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<true>));
+        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<true>, cudaFuncAttributeMaxDynamicSharedMemorySize,
+            static_cast<int>(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<true>, dsa_prefill::kThreads,
+            dsa_prefill::kInt8SplitKSharedBytes));
+        int8_tc_decode_resident_ctas = multiprocessors * blocks_per_multiprocessor;
+      }
     }
   }
   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,
+                                        int8_tc_decode_resident_ctas)
+            : workspace_splits;
     const int64_t partial_blocks = query_blocks * split_count;
     const int64_t grouped_blocks =
         static_cast<int64_t>(parameters.batch_size) * parameters.sequence_length * parameters.kv_num_heads *
@@ -2116,9 +2217,27 @@ 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<T, __nv_bfloat16>) {
+        const dim3 grid(static_cast<unsigned>(parameters.batch_size * parameters.sequence_length),
+                        static_cast<unsigned>(parameters.kv_num_heads), static_cast<unsigned>(split_count));
+        dsa_prefill::DynamicSparseAttentionInt8TcKernel<true>
+            <<<grid, dsa_prefill::kThreads, dsa_prefill::kInt8SplitKSharedBytes, stream>>>(
+                data.prepared_query, reinterpret_cast<const int8_t*>(data.present_key),
+                reinterpret_cast<const int8_t*>(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<int>("ORT_DSA_INT8_VECTOR_STAGING", 1);
       ORT_ENFORCE(vectorized_staging == 0 || vectorized_staging == 1,
@@ -2249,7 +2368,18 @@ Status LaunchDynamicSparseAttention(
     if constexpr (std::is_same_v<T, __nv_bfloat16>) {
       const dim3 grid(static_cast<unsigned>(parameters.batch_size * parameters.sequence_length),
                       static_cast<unsigned>(parameters.kv_num_heads));
-      if (data.quantized_cache) {
+      if (use_int8_tc_prefill) {
+        CUDA_RETURN_IF_ERROR(cudaFuncSetAttribute(
+            dsa_prefill::DynamicSparseAttentionInt8TcKernel<false>, cudaFuncAttributeMaxDynamicSharedMemorySize,
+            static_cast<int>(dsa_prefill::kInt8SharedBytes)));
+        dsa_prefill::DynamicSparseAttentionInt8TcKernel<false>
+            <<<grid, dsa_prefill::kThreads, dsa_prefill::kInt8SharedBytes, stream>>>(
+                data.prepared_query, reinterpret_cast<const int8_t*>(data.present_key),
+                reinterpret_cast<const int8_t*>(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<true>, cudaFuncAttributeMaxDynamicSharedMemorySize,
             static_cast<int>(dsa_prefill::kSharedBytes)));

onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h

@@ -50,6 +50,10 @@ struct DynamicSparseAttentionData {
   T* output = nullptr;
   float* attention_workspace = nullptr;
   bool quantized_cache = false;
+  bool prefill_kernel = 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;
 };