diff --git a/crates/onnx-runtime-ep-cuda/src/kernels/pointwise.rs b/crates/onnx-runtime-ep-cuda/src/kernels/pointwise.rs index 8402b93f8d..ed2b23f0f8 100644 --- a/crates/onnx-runtime-ep-cuda/src/kernels/pointwise.rs +++ b/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(sqrtf(load_float(x[i]))); \ + y[i] = store_float(1.0f / load_float(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]) -> Result> { + fn create(&self, node: &Node, _input_shapes: &[Vec]) -> Result> { + // `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(), })) } diff --git a/crates/onnx-runtime-ep-cuda/src/kernels/reduce.rs b/crates/onnx-runtime-ep-cuda/src/kernels/reduce.rs index 27662bb172..6fd05711da 100644 --- a/crates/onnx-runtime-ep-cuda/src/kernels/reduce.rs +++ b/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 diff --git a/crates/onnx-runtime-ep-cuda/src/optimizer.rs b/crates/onnx-runtime-ep-cuda/src/optimizer.rs index 74298f96a6..5a6b6336bf 100644 --- a/crates/onnx-runtime-ep-cuda/src/optimizer.rs +++ b/crates/onnx-runtime-ep-cuda/src/optimizer.rs @@ -8,6 +8,17 @@ use onnx_runtime_optimizer::{ pub(crate) const SILU_MUL_FUSION_ATTR: &str = "_cuda_silu_mul"; pub(crate) const DECOMPOSED_SILU_ATTR: &str = "_cuda_decomposed_silu"; +/// Private marker set by [`CudaRsqrtFusion`] on a `Reciprocal` node whose input +/// was a single-consumer `Sqrt`. The fusion rewires the `Reciprocal` to read the +/// pre-`Sqrt` value and deletes the `Sqrt`; the CUDA unary-math factory +/// ([`crate::kernels::pointwise::UnaryMathFactory`]) sees this marker and +/// dispatches the single-kernel `rsqrt_{dtype}` (byte-identical double-round) +/// instead of the plain `reciprocal_{dtype}`. The node keeps op_type +/// `Reciprocal` in the default domain so standard unary shape inference and +/// registry dispatch are unchanged, and the session restores the pre-pass graph +/// before any non-CUDA fallback so no other EP observes the rewired node. +pub(crate) const CUDA_RSQRT_ATTR: &str = "_cuda_rsqrt"; + /// Private marker set on a `MatMulNBits` node whose trailing bias input came /// from folding a *separate* elementwise `Add(MatMulNBits(x), bias)`. /// @@ -90,6 +101,11 @@ pub(crate) struct CudaSiluFusion; pub(crate) fn cuda_optimization_passes() -> Vec> { vec![ Box::new(CudaSiluFusion), + // Collapse the SSM/linear-attention `Reciprocal(Sqrt(x))` normalize-scale + // glue into a single byte-identical `rsqrt` kernel (drops the standalone + // reciprocal kernel + its fp16 intermediate round-trip, ~1.2% of decode + // GPU-kernel time on qwen3.5-0.8b-hybrid; see the profiling drop). + Box::new(CudaRsqrtFusion), Box::new(CudaFoldConstantTranspose), Box::new(CudaFoldConstantCast), // Fuse the per-layer Q/K/V projections into one wider MatMulNBits while @@ -202,6 +218,110 @@ impl OptimizationPass for CudaSiluFusion { } } +/// Collapse an ONNX `Reciprocal(Sqrt(x))` pair into a single fused `rsqrt` +/// kernel. +/// +/// ## The pattern +/// +/// SSM / linear-attention decoders (e.g. qwen3.5-*-hybrid) normalize the SSM +/// query/key by a sum-of-squares scale, which exporters emit as the primitive +/// chain `... Add(sumsq, eps) → Sqrt → Reciprocal → Mul(x, scale)`. The `Sqrt` +/// and `Reciprocal` land as two standalone per-layer kernels +/// (`sqrt_{dtype}` + `reciprocal_{dtype}`); a captured-graph node trace of +/// qwen3.5-0.8b-hybrid decode measured the standalone `reciprocal_f16` at ~1.2% +/// of decode GPU-kernel time (fires 2× per SSM layer), each reading and writing +/// the full scale tensor. +/// +/// ## The rewrite +/// +/// For every `Sqrt` whose single output is consumed *only* by one `Reciprocal` +/// (and is not a graph output), this pass rewires the `Reciprocal` to read the +/// `Sqrt`'s input directly, tags it with [`CUDA_RSQRT_ATTR`], and deletes the +/// `Sqrt`. The node keeps op_type `Reciprocal` in the default domain, so +/// standard unary shape inference and kernel-registry dispatch are unchanged; +/// the CUDA unary-math factory reads the marker and launches the fused +/// `rsqrt_{dtype}` kernel, which reproduces the exact two-kernel rounding +/// (`Sqrt` rounds `sqrtf(x)` to the storage dtype, then `Reciprocal` reads that +/// back) so greedy tokens stay byte-identical while the intermediate tensor's +/// global round-trip and one kernel launch are eliminated. +/// +/// This is model-agnostic: it matches the structural `Sqrt → Reciprocal` glue +/// on any shape/dtype the unary kernels support (f32/f16/bf16), not a single +/// architecture. The single-consumer / no-escape guard guarantees the fusion +/// never removes a `Sqrt` value another node still observes. +#[derive(Clone, Copy, Debug, Default)] +pub(crate) struct CudaRsqrtFusion; + +impl OptimizationPass for CudaRsqrtFusion { + fn name(&self) -> &str { + "CudaRsqrtFusion" + } + + fn run(&self, graph: &mut Graph, _ctx: &PassContext) -> OptimizerResult<()> { + let sqrt_ids: Vec = graph + .nodes + .iter() + .filter_map(|(id, node)| { + (node.op_type == "Sqrt" + && node.is_default_domain() + && node.inputs.len() == 1 + && node.outputs.len() == 1) + .then_some(id) + }) + .collect(); + let mut changed = false; + + for sqrt_id in sqrt_ids { + let Some(sqrt) = graph.try_node(sqrt_id) else { + continue; + }; + let Some(x) = sqrt.inputs[0] else { + continue; + }; + // The fused `rsqrt` kernel exists for f32/f16/bf16 and reproduces the + // two-op rounding exactly, so restrict the fusion to those dtypes. + if !matches!( + graph.value(x).dtype, + DataType::Float32 | DataType::Float16 | DataType::BFloat16 + ) { + continue; + } + let sqrt_output = sqrt.outputs[0]; + if graph.outputs.contains(&sqrt_output) { + continue; + } + let consumers = graph.consumers(sqrt_output); + if consumers.len() != 1 { + continue; + } + let recip_id = consumers[0]; + let recip = graph.node(recip_id); + if recip.op_type != "Reciprocal" + || !recip.is_default_domain() + || recip.inputs.len() != 1 + || recip.outputs.len() != 1 + || recip.inputs[0] != Some(sqrt_output) + || recip.attr(CUDA_RSQRT_ATTR).is_some() + { + continue; + } + + let mut rsqrt = recip.clone(); + rsqrt.inputs = vec![Some(x)]; + rsqrt.attributes + .insert(CUDA_RSQRT_ATTR.into(), Attribute::Int(1)); + graph.replace_node(recip_id, rsqrt); + graph.remove_node(sqrt_id); + changed = true; + } + + if changed { + graph.validate().map_err(OptimizerError::from)?; + } + Ok(()) + } +} + /// Drop the redundant `Cast` pairs that some exporters (e.g. Phi-4-mini, /// Muse-Glimmer) wrap around every simplified-layer-norm / RMS-norm, running the /// norm directly on its narrow (fp16 or bf16) activations instead. @@ -3318,6 +3438,100 @@ mod tests { ); } + // === Rsqrt fusion (Reciprocal(Sqrt(x)) -> tagged Reciprocal) === + + /// `Add(sumsq, eps) → Sqrt → Reciprocal → Mul(x, scale)`, the SSM + /// normalize-scale glue `CudaRsqrtFusion` targets. The `Add` and trailing + /// `Mul` bracket the pair so the single-consumer rewrite is observable. + /// Returns the graph plus the `denom` (pre-Sqrt) and `root` (Sqrt output) + /// value ids for assertions. + fn rsqrt_glue_graph(dtype: DataType) -> (Graph, ValueId, ValueId) { + let mut graph = Graph::new(); + graph.opset_imports.insert(String::new(), 17); + let sumsq = value(&mut graph, "sumsq", dtype, 5); + let eps = value(&mut graph, "eps", dtype, 5); + let x = value(&mut graph, "x", dtype, 5); + let denom = value(&mut graph, "denom", dtype, 5); + let root = value(&mut graph, "root", dtype, 5); + let scale = value(&mut graph, "scale", dtype, 5); + let output = value(&mut graph, "output", dtype, 5); + graph.add_input(sumsq); + graph.add_input(eps); + graph.add_input(x); + graph.insert_node(Node::new( + NodeId(0), + "Add", + vec![Some(sumsq), Some(eps)], + vec![denom], + )); + graph.insert_node(Node::new(NodeId(0), "Sqrt", vec![Some(denom)], vec![root])); + graph.insert_node(Node::new( + NodeId(0), + "Reciprocal", + vec![Some(root)], + vec![scale], + )); + graph.insert_node(Node::new( + NodeId(0), + "Mul", + vec![Some(x), Some(scale)], + vec![output], + )); + graph.add_output(output); + (graph, denom, root) + } + + #[test] + fn fuses_fp16_reciprocal_of_sqrt() { + let (mut graph, denom, _root) = rsqrt_glue_graph(DataType::Float16); + CudaRsqrtFusion.run(&mut graph, &PassContext::new()).unwrap(); + + // Sqrt removed; Reciprocal retagged and rewired to the pre-sqrt value. + assert_eq!(graph.num_nodes(), 3, "Sqrt must be removed"); + assert!( + graph.nodes.values().all(|n| n.op_type != "Sqrt"), + "Sqrt node must be gone" + ); + let recip = graph + .nodes + .values() + .find(|n| n.op_type == "Reciprocal") + .expect("Reciprocal must remain"); + assert!(recip.is_default_domain(), "node stays a standard Reciprocal"); + assert_eq!( + recip.attr(CUDA_RSQRT_ATTR).and_then(Attribute::as_int), + Some(1), + "fused Reciprocal must carry the rsqrt marker" + ); + // Reciprocal now reads the pre-Sqrt value ("denom"), not the Sqrt output. + assert_eq!(recip.inputs[0], Some(denom)); + assert!(graph.validate().is_ok()); + } + + #[test] + fn does_not_fuse_sqrt_with_extra_consumer() { + // Sqrt output feeds both the Reciprocal *and* a second consumer, so the + // Sqrt value escapes the pair and must not be removed. + let (mut graph, _denom, root) = rsqrt_glue_graph(DataType::Float16); + let sink = value(&mut graph, "sink", DataType::Float16, 5); + graph.insert_node(Node::new(NodeId(0), "Neg", vec![Some(root)], vec![sink])); + graph.add_output(sink); + + CudaRsqrtFusion.run(&mut graph, &PassContext::new()).unwrap(); + assert!( + graph.nodes.values().any(|n| n.op_type == "Sqrt"), + "Sqrt with a second consumer must be preserved" + ); + assert!( + graph + .nodes + .values() + .filter(|n| n.op_type == "Reciprocal") + .all(|n| n.attr(CUDA_RSQRT_ATTR).is_none()), + "Reciprocal must not be retagged when Sqrt escapes" + ); + } + // === QKV bias fold === use onnx_runtime_ir::{TensorData, WeightRef}; diff --git a/crates/onnx-runtime-ep-cuda/tests/pointwise_gpu.rs b/crates/onnx-runtime-ep-cuda/tests/pointwise_gpu.rs index c68be17f9f..17dc97baa0 100644 --- a/crates/onnx-runtime-ep-cuda/tests/pointwise_gpu.rs +++ b/crates/onnx-runtime-ep-cuda/tests/pointwise_gpu.rs @@ -20,7 +20,9 @@ use onnx_runtime_ep_api::{ }; use onnx_runtime_ep_cuda::CudaExecutionProvider; use onnx_runtime_ep_cuda::runtime::cuptr; -use onnx_runtime_ir::{DataType, Graph, Node, NodeId, compute_contiguous_strides, static_shape}; +use onnx_runtime_ir::{ + Attribute, DataType, Graph, Node, NodeId, compute_contiguous_strides, static_shape, +}; use onnx_runtime_loader::Model; fn encode(values: &[f32], dtype: DataType) -> Vec { @@ -147,6 +149,43 @@ fn run_unary(ep: &CudaExecutionProvider, op: &str, dtype: DataType, x: &[f32]) - decode(&out, dtype) } +/// Run the fused `rsqrt` path: a standard-domain `Reciprocal` node carrying the +/// `_cuda_rsqrt` marker (as `CudaRsqrtFusion` produces), which the CUDA +/// unary-math factory dispatches to the single-kernel `rsqrt_{dtype}`. +fn run_rsqrt(ep: &CudaExecutionProvider, dtype: DataType, x: &[f32]) -> Vec { + let rt = ep.runtime(); + let bytes = encode(x, dtype); + let x_buf = ep.allocate(bytes.len(), 256).unwrap(); + let mut y_buf = ep.allocate(bytes.len(), 256).unwrap(); + unsafe { rt.htod(&bytes, cuptr(x_buf.as_ptr())).unwrap() }; + let shape = [x.len()]; + let strides = compute_contiguous_strides(&shape); + let input = TensorView::new( + DevicePtr(x_buf.as_ptr()), + dtype, + &shape, + &strides, + ep.device_id(), + ); + let output = TensorMut::new( + DevicePtrMut(y_buf.as_mut_ptr()), + dtype, + &shape, + &strides, + ep.device_id(), + ); + let mut node = Node::new(NodeId(0), "Reciprocal", vec![], vec![]); + node.attributes + .insert("_cuda_rsqrt".into(), Attribute::Int(1)); + let kernel = ep.get_kernel(&node, &[], 17).unwrap(); + kernel.execute(&[input], &mut [output]).unwrap(); + let mut out = vec![0u8; bytes.len()]; + unsafe { rt.dtoh(&mut out, cuptr(y_buf.as_ptr())).unwrap() }; + ep.deallocate(x_buf).unwrap(); + ep.deallocate(y_buf).unwrap(); + decode(&out, dtype) +} + fn run_predicate( ep: &CudaExecutionProvider, op: &str, @@ -430,6 +469,37 @@ fn half_unary_and_activation_families_match_cpu_reference() { } } +/// The fused `rsqrt` kernel (from `CudaRsqrtFusion`) must be **byte-identical** +/// to the unfused `Reciprocal(Sqrt(x))` two-kernel path — that is the parity +/// contract that keeps native decode greedy tokens unchanged. It reproduces the +/// intermediate narrow-dtype rounding (`Sqrt` rounds `sqrtf(x)` to the storage +/// dtype, then `Reciprocal` reads that back), so an exact-equality check (not a +/// tolerance) is required. +#[cfg_attr( + not(feature = "gpu-tests"), + ignore = "requires CUDA device; enable the gpu-tests feature on a CUDA runner" +)] +#[test] +fn fused_rsqrt_is_byte_identical_to_sqrt_then_reciprocal() { + let ep = require_cuda(); + // Positive inputs across a wide magnitude range (rsqrt domain); avoids NaN + // so exact equality is well-defined. + let x = [ + 1.0e-4, 0.01, 0.25, 0.5, 1.0, 2.0, 3.0, 7.5, 42.0, 123.5, 1000.0, 65504.0, + ]; + for dtype in [DataType::Float16, DataType::BFloat16] { + // Two-kernel reference: Sqrt then Reciprocal, each rounding to `dtype`. + let sqrted = run_unary(&ep, "Sqrt", dtype, &x); + let reference = run_unary(&ep, "Reciprocal", dtype, &sqrted); + let fused = run_rsqrt(&ep, dtype, &x); + assert_eq!( + fused.iter().map(|v| v.to_bits()).collect::>(), + reference.iter().map(|v| v.to_bits()).collect::>(), + "fused rsqrt must be bit-identical to Reciprocal(Sqrt(x)) for {dtype:?}" + ); + } +} + #[cfg_attr( not(feature = "gpu-tests"), ignore = "requires CUDA device; enable the gpu-tests feature on a CUDA runner"