Fuse SSM/linear-attention reduction: drop fp32 cuDNN ReduceSum round-trip on native decode

crates/onnx-runtime-ep-cuda/src/kernels/pointwise.rs

@@ -36,7 +36,7 @@ use std::sync::{Arc, Mutex};
 use cudarc::driver::{LaunchConfig, PushKernelArg};
 
 use onnx_runtime_ep_api::{EpError, Kernel, KernelFactory, Result, TensorMut, TensorView};
-use onnx_runtime_ir::{DataType, Node};
+use onnx_runtime_ir::{Attribute, DataType, Node};
 
 use super::elementwise::{
     BroadcastMetadataCache, BroadcastMetadataKey, capture_shape_eligible,
@@ -193,14 +193,33 @@ DEFINE_UNARY(atan, TYPE, SUFFIX) \
 DEFINE_UNARY(asinh, TYPE, SUFFIX) \
 DEFINE_UNARY(acosh, TYPE, SUFFIX) \
 DEFINE_UNARY(atanh, TYPE, SUFFIX)
+// Fused reciprocal-square-root: collapses an ONNX `Reciprocal(Sqrt(x))` glue
+// pair (see `CudaRsqrtFusion`) into one kernel. It deliberately reproduces the
+// exact two-kernel rounding — `Sqrt` first rounds `sqrtf(x)` to the storage
+// dtype, then `Reciprocal` reads that rounded intermediate back — so the fused
+// result is byte-identical to the unfused pair (critical for byte-identical
+// greedy decode), while eliminating the intermediate tensor's global round-trip
+// and one kernel launch. Correct for f32/f16/bf16 (for f32 the intermediate
+// round is a no-op, so it matches the two-op f32 path exactly).
+#define DEFINE_RSQRT(TYPE, SUFFIX) \
+extern "C" __global__ void rsqrt_##SUFFIX(const TYPE* x, TYPE* y, const unsigned long long n) { \
+    for (unsigned long long i = blockIdx.x*blockDim.x + threadIdx.x; i < n; \
+         i += (unsigned long long)gridDim.x * blockDim.x) { \
+        const TYPE s = store_float<TYPE>(sqrtf(load_float<TYPE>(x[i]))); \
+        y[i] = store_float<TYPE>(1.0f / load_float<TYPE>(s)); \
+    } \
+}
 DEFINE_FOR_TYPE(float, f32)
+DEFINE_RSQRT(float, f32)
 #ifdef NXRT_HAS_CUDA_HALF_HEADERS
 DEFINE_FOR_TYPE(__half, f16)
 DEFINE_FOR_TYPE(__nv_bfloat16, bf16)
+DEFINE_RSQRT(__half, f16)
+DEFINE_RSQRT(__nv_bfloat16, bf16)
 #endif
 "#;
 
-const UNARY_MATH_MODULE: &str = "pointwise_unary_math_float_v3";
+const UNARY_MATH_MODULE: &str = "pointwise_unary_math_float_v4";
 
 /// A supported unary math op and its NVRTC entry point.
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
@@ -226,6 +245,10 @@ pub enum UnaryMathOp {
     Asinh,
     Acosh,
     Atanh,
+    /// Fused `Reciprocal(Sqrt(x))`. Not a standard ONNX op — only produced by
+    /// [`crate::optimizer::CudaRsqrtFusion`], which retags a `Reciprocal` node
+    /// whose input is a single-consumer `Sqrt`.
+    Rsqrt,
 }
 
 impl UnaryMathOp {
@@ -252,6 +275,7 @@ impl UnaryMathOp {
             UnaryMathOp::Asinh => "asinh",
             UnaryMathOp::Acosh => "acosh",
             UnaryMathOp::Atanh => "atanh",
+            UnaryMathOp::Rsqrt => "rsqrt",
         }
     }
 
@@ -282,6 +306,9 @@ impl UnaryMathOp {
             UnaryMathOp::Asinh => "Asinh",
             UnaryMathOp::Acosh => "Acosh",
             UnaryMathOp::Atanh => "Atanh",
+            // Reported in errors as the source op; the fused node keeps op_type
+            // `Reciprocal`.
+            UnaryMathOp::Rsqrt => "Reciprocal",
         }
     }
 }
@@ -293,9 +320,22 @@ pub struct UnaryMathFactory {
 }
 
 impl KernelFactory for UnaryMathFactory {
-    fn create(&self, _node: &Node, _input_shapes: &[Vec<usize>]) -> Result<Box<dyn Kernel>> {
+    fn create(&self, node: &Node, _input_shapes: &[Vec<usize>]) -> Result<Box<dyn Kernel>> {
+        // `CudaRsqrtFusion` retags a `Reciprocal` whose input is a
+        // single-consumer `Sqrt` with the `_cuda_rsqrt` marker and rewires it to
+        // the pre-`sqrt` value, so this node computes `1/sqrt(x)` in one kernel.
+        let op = if self.op == UnaryMathOp::Reciprocal
+            && node
+                .attr(crate::optimizer::CUDA_RSQRT_ATTR)
+                .and_then(Attribute::as_int)
+                == Some(1)
+        {
+            UnaryMathOp::Rsqrt
+        } else {
+            self.op
+        };
         Ok(Box::new(UnaryMathKernel {
-            op: self.op,
+            op,
             runtime: self.runtime.clone(),
         }))
     }

crates/onnx-runtime-ep-cuda/src/kernels/reduce.rs

@@ -2,10 +2,13 @@
 //! (`docs/execution/CUDA_COVERAGE.md`, "Normalization & softmax" / reduce rows).
 //!
 //! `ReduceSum` and `ReduceMean` use `cudnnReduceTensor` for f32 (native f32
-//! comp type) and f16 (with an f32 comp type). bf16 cannot be reduced by cuDNN
-//! (rejected even with an f32 comp type — `CUDNN_STATUS_NOT_SUPPORTED`), so it
-//! uses the typed NVRTC block reduction, which accumulates in f32. That same
-//! NVRTC fallback also runs for f32/f16 when cuDNN is absent.
+//! comp type). f16 and bf16 use the typed NVRTC block reduction, which does
+//! fp16/bf16 IO with f32 register accumulation in a single capture-safe kernel.
+//! For f16 this deliberately avoids cuDNN's fp16↔fp32 `op_tensor` cast +
+//! fp32-ReduceSum round-trip (cuDNN rejects a half `reduceTensorCompType` and
+//! forces `CUDNN_DATA_FLOAT`); bf16 cannot be reduced by cuDNN at all (rejected
+//! even with an f32 comp type — `CUDNN_STATUS_NOT_SUPPORTED`). That same NVRTC
+//! path also runs for f32 when cuDNN is absent.
 //! `ReduceMax`/`ReduceMin` always use NVRTC.
 //!
 //! `cub::DeviceReduce` / `DeviceSegmentedReduce` are the vendor primitives for
@@ -1033,13 +1036,26 @@ impl ReduceKernel {
             return Ok(());
         }
 
-        // cuDNN reduces f32 (native comp type) and f16 (with an f32 comp type,
-        // see `CudnnBackend::reduce`). It cannot reduce bf16 at all — cuDNN
-        // rejects a bf16 tensor even with an f32 comp type
-        // (`CUDNN_STATUS_NOT_SUPPORTED`, reproduced on cuDNN 9.10 and 9.20) — so
-        // bf16 falls through to the typed NVRTC block reduction below, which
-        // accumulates in f32 (identical ONNX semantics).
-        if matches!(x.dtype, DataType::Float32 | DataType::Float16)
+        // cuDNN reduces f32 with a native f32 comp type — byte-identical to the
+        // typed NVRTC block reduction but on cuDNN's tuned kernels — so f32 sum/
+        // mean stays here. f16 does **not**: `cudnnReduceTensor` rejects a half
+        // `reduceTensorCompType` and forces `CUDNN_DATA_FLOAT`, which cuDNN
+        // implements as an fp16→fp32 `op_tensor` cast, an fp32 ReduceSum, and an
+        // fp32→fp16 cast — a three-kernel round-trip through a full-size fp32
+        // temporary. A captured-graph node trace of the qwen3.5-0.8b-hybrid
+        // decode (Deckard, `.squad/decisions/inbox/deckard-fair-hybrid-gap.md`)
+        // measured that fragmentation at ~20.7% of decode GPU-kernel time on the
+        // SSM/linear-attention path (cudnn ReduceSum 8.1% + fp16↔fp32 op_tensor
+        // casts 2.7%, firing 18× per step). The typed NVRTC block reduction
+        // below already does fp16 IO with fp32 register accumulation — the exact
+        // ONNX accumulate-in-f32-then-cast rule — in a **single** capture-safe
+        // kernel, with no fp32 temporary and half the memory traffic, so f16
+        // sum/mean is routed there instead. This is general across all shapes/
+        // axes/dims (the same path already serves bf16 and every extended f16
+        // reduction), not a per-model special case. bf16 cannot be reduced by
+        // cuDNN at all (`CUDNN_STATUS_NOT_SUPPORTED`, cuDNN 9.10/9.20) and also
+        // falls through below.
+        if x.dtype == DataType::Float32
             && let Some(cudnn_op) = cudnn_op
             && self.runtime.cudnn().is_available()
         {
@@ -1086,9 +1102,11 @@ impl ReduceKernel {
             return Ok(());
         }
 
-        // NVRTC block-reduction path (Int64 DATA reduce; bf16, which cuDNN
-        // cannot reduce; every extended op — `ReduceSumSquare`/L1/L2/Prod/
-        // LogSum(Exp) — routed by `ext_tags`; and f16/f32 when cuDNN is absent).
+        // NVRTC block-reduction path (Int64 DATA reduce; f16 and bf16 sum/mean,
+        // which take the single-kernel fused fp16/bf16-IO f32-accumulation path
+        // here rather than cuDNN's fp32 round-trip; every extended op —
+        // `ReduceSumSquare`/L1/L2/Prod/LogSum(Exp) — routed by `ext_tags`; and
+        // f32 when cuDNN is absent).
         // All of these use the i64 base/delta offset tables, so they share one
         // capture-eligible metadata cache: a shape-stable decode reduce reuses
         // the cached device tables with no per-call `alloc`/`htod`/`free`, gates