perf(ArcAttention): port vLLM rms_norm_kernel — fuse the V4 per-head Q RMSNorm (7 launches -> 1)

mistralrs-core/build.rs

@@ -1,6 +1,79 @@
 #[cfg(feature = "cuda")]
 const CUDA_NVCC_FLAGS: Option<&'static str> = option_env!("CUDA_NVCC_FLAGS");
 
+/// CUDA sources whose output must be BIT-IDENTICAL to candle-kernels, and
+/// which therefore must be compiled without `--use_fast_math`. Each carries a
+/// bit-identity contract in its file header.
+#[cfg(feature = "cuda")]
+const IEEE_SOURCES: &[&str] = &["sinkhorn.cu", "swiglu_clamp.cu", "qnorm.cu"];
+
+/// nvcc args for the normal (fast-math) `src/cuda/*.cu` builder.
+#[cfg(feature = "cuda")]
+const FAST_MATH_ARGS: &[&str] = &[
+    "-std=c++17",
+    "-O3",
+    "-U__CUDA_NO_HALF_OPERATORS__",
+    "-U__CUDA_NO_HALF_CONVERSIONS__",
+    "-U__CUDA_NO_HALF2_OPERATORS__",
+    "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
+    "--expt-relaxed-constexpr",
+    "--expt-extended-lambda",
+    "--use_fast_math",
+    "--verbose",
+];
+
+/// nvcc args for the IEEE (no fast-math) builder.
+#[cfg(feature = "cuda")]
+const IEEE_ARGS: &[&str] = &[
+    "-std=c++17",
+    "-O3",
+    "--expt-relaxed-constexpr",
+    "--expt-extended-lambda",
+    "--fmad=false",
+    "--verbose",
+];
+
+/// The guard that the `#error` idiom cannot be: nvcc 12.4 defines NO
+/// preprocessor macro for `--use_fast_math` in either compilation pass (probed
+/// directly on sm_90 with this toolchain), so a `.cu`-side `#if
+/// defined(__USE_FAST_MATH__)` — as sinkhorn.cu carries today — never fires.
+///
+/// This asserts the invariant at the only place it can regress: the flag sets
+/// in this file. Adding `--use_fast_math` to `IEEE_ARGS`, or dropping a source
+/// from the fast-math builder's exclusion list, becomes a hard build failure
+/// instead of silent numeric drift in generated tokens.
+///
+/// Proved red: temporarily appending `"--use_fast_math"` to `IEEE_ARGS` makes
+/// `cargo build -p mistralrs-core --features cuda` panic here before nvcc is
+/// ever invoked.
+#[cfg(feature = "cuda")]
+fn assert_ieee_kernel_flags(sources: &[&str], fast_math_args: &[&str], ieee_args: &[&str]) {
+    assert!(
+        fast_math_args.contains(&"--use_fast_math"),
+        "the general src/cuda/*.cu builder lost --use_fast_math; if that is \
+         intentional, IEEE_SOURCES no longer needs a separate builder"
+    );
+    assert!(
+        !ieee_args.contains(&"--use_fast_math"),
+        "BIT-IDENTITY VIOLATION: the IEEE builder was given --use_fast_math. \
+         {sources:?} must match candle-kernels bit for bit, and fast math \
+         rewrites expf -> ex2.approx.ftz.f32 and div.rn.f32 -> \
+         div.approx.ftz.f32 (measured, nvcc 12.4/sm_90). This changes \
+         generated tokens."
+    );
+    for src in sources {
+        let path = std::path::Path::new("src/cuda").join(src);
+        assert!(
+            path.exists(),
+            "IEEE_SOURCES lists {src}, which does not exist at {}. It would be \
+             silently dropped from the build while still being excluded from \
+             the fast-math glob.",
+            path.display()
+        );
+        println!("cargo:rerun-if-changed={}", path.display());
+    }
+}
+
 fn main() {
     set_git_revision();
 
@@ -10,28 +83,29 @@ fn main() {
         println!("cargo:rerun-if-changed=build.rs");
         let build_dir = PathBuf::from(std::env::var("OUT_DIR").unwrap());
 
-        // sinkhorn.cu is EXCLUDED from this fast-math builder and compiled
-        // separately below: it must be bit-identical to candle-kernels (which
-        // build with plain -O3, no fast math), and --use_fast_math rewrites
-        // expf -> __expf and IEEE division -> approximate reciprocals. The
-        // kernel source carries an `#error` guard against being re-globbed
-        // under fast math. See mistralrs-core/src/cuda/sinkhorn.cu.
+        // These sources are EXCLUDED from this fast-math builder and compiled
+        // separately below: they must be bit-identical to candle-kernels
+        // (which build with plain -O3, no fast math), and --use_fast_math
+        // rewrites expf -> __expf and IEEE division -> approximate
+        // reciprocals. See mistralrs-core/src/cuda/{sinkhorn,swiglu_clamp}.cu.
+        //
+        // NOTE: the `#if defined(__USE_FAST_MATH__)` `#error` in sinkhorn.cu is
+        // VACUOUS. nvcc 12.4 defines no preprocessor macro for
+        // --use_fast_math in either the host or the device pass (probed on
+        // sm_90: __USE_FAST_MATH__, __CUDA_FAST_MATH__, __FAST_MATH__,
+        // __CUDACC_FAST_MATH__, __CUDA_PREC_DIV__, __CUDA_FTZ__ are all
+        // undefined with and without the flag), so that guard can never fire.
+        // `assert_ieee_kernel_flags` below is the guard that does.
+        assert_ieee_kernel_flags(IEEE_SOURCES, FAST_MATH_ARGS, IEEE_ARGS);
+
         let mut builder = cudaforge::KernelBuilder::new()
             .source_glob("src/cuda/*.cu")
-            .exclude(&["sinkhorn.cu"])
-            .out_dir(&build_dir)
-            .arg("-std=c++17")
-            .arg("-O3")
-            .arg("-U__CUDA_NO_HALF_OPERATORS__")
-            .arg("-U__CUDA_NO_HALF_CONVERSIONS__")
-            .arg("-U__CUDA_NO_HALF2_OPERATORS__")
-            .arg("-U__CUDA_NO_BFLOAT16_CONVERSIONS__")
-            .arg("--expt-relaxed-constexpr")
-            .arg("--expt-extended-lambda")
-            .arg("--use_fast_math")
-            .arg("--verbose")
-            .arg("--compiler-options")
-            .arg("-fPIC");
+            .exclude(IEEE_SOURCES)
+            .out_dir(&build_dir);
+        for a in FAST_MATH_ARGS {
+            builder = builder.arg(a);
+        }
+        builder = builder.arg("--compiler-options").arg("-fPIC");
 
         // Check if CUDA_COMPUTE_CAP < 80 and disable bf16 kernels if so.
         // bf16 WMMA operations and certain bf16 intrinsics are only available on sm_80+.
@@ -63,23 +137,27 @@ fn main() {
         println!("cargo:rustc-link-search={}", build_dir.display());
         println!("cargo:rustc-link-lib=mistralrscuda");
 
-        // Dedicated IEEE (no fast math) builder for sinkhorn.cu — bit-identity
-        // with candle-kernels requires accurate expf + div.rn.f32; --fmad=false
-        // additionally forbids FMA contraction so rounding matches candle's
-        // unfused op chain exactly. Own subdirectory so its build cache never
+        // Dedicated IEEE (no fast math) builder for the bit-identity kernels —
+        // bit-identity with candle-kernels requires accurate expf +
+        // div.rn.f32; --fmad=false additionally forbids FMA contraction so
+        // rounding matches candle's unfused op chain exactly. (Verified on
+        // nvcc 12.4/sm_90 that --fmad=false does NOT perturb libdevice's
+        // accurate expf expansion, which uses explicit fma.rn.f32 rather than
+        // source-level contraction.) Own subdirectory so its build cache never
         // mixes with the fast-math builder's.
         let sinkhorn_dir = build_dir.join("sinkhorn_ieee");
         let mut sinkhorn_builder = cudaforge::KernelBuilder::new()
-            .source_files(vec!["src/cuda/sinkhorn.cu"])
-            .out_dir(&sinkhorn_dir)
-            .arg("-std=c++17")
-            .arg("-O3")
-            .arg("--expt-relaxed-constexpr")
-            .arg("--expt-extended-lambda")
-            .arg("--fmad=false")
-            .arg("--verbose")
-            .arg("--compiler-options")
-            .arg("-fPIC");
+            .source_files(
+                IEEE_SOURCES
+                    .iter()
+                    .map(|f| format!("src/cuda/{f}"))
+                    .collect::<Vec<_>>(),
+            )
+            .out_dir(&sinkhorn_dir);
+        for a in IEEE_ARGS {
+            sinkhorn_builder = sinkhorn_builder.arg(a);
+        }
+        sinkhorn_builder = sinkhorn_builder.arg("--compiler-options").arg("-fPIC");
         if let Some(cuda_nvcc_flags_env) = CUDA_NVCC_FLAGS {
             sinkhorn_builder = sinkhorn_builder.arg("--compiler-options");
             sinkhorn_builder = sinkhorn_builder.arg(cuda_nvcc_flags_env);

mistralrs-core/src/cuda/ffi.rs

@@ -371,4 +371,43 @@ extern "C" {
         eps: f32,
         stream: i64,
     );
+    // Fused clamped-SwiGLU (cuda/swiglu_clamp.cu). Compiled by the dedicated
+    // IEEE (no fast-math) builder in build.rs — see the bit-identity contract
+    // in that file's header.
+    pub(crate) fn swiglu_clamp_bf16(
+        gate: *const c_void,
+        up: *const c_void,
+        out: *mut c_void,
+        n: u32,
+        limit: f32,
+        stream: i64,
+    );
+    pub(crate) fn swiglu_clamp_f16(
+        gate: *const c_void,
+        up: *const c_void,
+        out: *mut c_void,
+        n: u32,
+        limit: f32,
+        stream: i64,
+    );
+    pub(crate) fn swiglu_clamp_f32(
+        gate: *const c_void,
+        up: *const c_void,
+        out: *mut c_void,
+        n: u32,
+        limit: f32,
+        stream: i64,
+    );
+    // Fused per-head Q RMSNorm (cuda/qnorm.cu). Also IEEE-built. Returns 0 on
+    // success, non-zero if the shape is outside the single-block regime.
+    pub(crate) fn qnorm_rms_bf16(
+        x: *const c_void,
+        out: *mut c_void,
+        rows: u32,
+        cols: u32,
+        inv_n_bits: u16,
+        eps_bits: u16,
+        f32_acc: i32,
+        stream: i64,
+    ) -> i32;
 }

mistralrs-core/src/cuda/mod.rs

@@ -1,5 +1,7 @@
 pub mod ffi;
 pub mod gdn;
 pub mod moe;
+pub mod qnorm;
 pub mod sinkhorn;
 pub mod ssm;
+pub mod swiglu_clamp;