From 66e07326c3376c636db39cc13e6071e58f7be46f Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Tue, 18 Aug 2026 10:35:30 +0100 Subject: [PATCH 01/27] feat(ArcKV/Fp8): fused E4M3 quantize+dequantize CUDA kernel, default path The V4 B=1 decode step pays 43 blocking `cuMemcpyDtoHAsync_v2` per token (~109us each, 4.81 ms/token) inside `dsv4_kv_fp8::e4m3_codes_cpu`, which round-trips every layer's scaled K block through the HOST because candle has no CUDA F8E4M3 cast. Those copies are invisible to the obvious grep (`*Synchronize*` = 0.0 calls/step) and they make CUDA graph capture impossible: a graph cannot record a blocking D2H. The measured trap: the existing sync-free path (`ARC_GPU_ACT_QUANT=1`, `GpuApprox`) is SLOWER on H200 - interleaved A1 66.99 / B1 67.71 / A2 68.05 / B2 70.30 ms/token, with `kv_fp8_quant` 73.49 -> 134.18 us/call (+83%). It swaps one blocking copy for ~19 extra elementwise launches per layer, and this machine is bottlenecked on op count (9,131 launches/token, median kernel 1.18us, memory controller 4% utilized), not kernel speed. So `GpuApprox` is not the fix and is documented here as not being one. This adds `mistralrs_quant::arc_kvquant`: one fused kernel per direction, replacing ~11 candle ops on the quantize side and ~13 on the dequantize side with a single launch each, and removing the D2H entirely. The byte format is unchanged. Bit-parity with the CPU path is the bar, and it is obtained by construction rather than by hope: * the E4M3 rounding is a transcription of NVIDIA's `__nv_cvt_double_to_fp8(x, SATFINITE, E4M3)`, which is what the Rust `float8` crate ports and therefore what `F8E4M3::from_f32` computes; * `scale` reproduces `(amax / 448.0)?.affine(1.0, 1e-12)` exactly, including that candle lowers `Tensor / f64` to a MULTIPLY by the f32-rounded reciprocal; * mistralrs-quant compiles with `--use_fast_math`, so every float op is an explicit `__f*_rn` intrinsic (IEEE, unaffected by -prec-div/-ftz) and the amax reduction runs on absolute-value bit patterns as unsigned integers rather than through `fabsf`/`fmaxf`; * dequant indexes the SAME 256-entry `F8E4M3::from_bits(i).to_f32()` table the candle path fed to `index_select`. D33: the kernel ships a deliberate mutant (RNE replaced by truncation, everything else identical) reachable only from the parity test, so the comparison is shown to fail on a wrong kernel. The test also asserts the fused call counters moved - a parity check on this exact subsystem has already passed vacuously by comparing an implementation to itself. D14: the GPU parity test exits 2 (environment failure) rather than passing when no CUDA device is present. Co-Authored-By: Claude Opus 5 (1M context) --- mistralrs-core/src/models/dsv4_kv_fp8.rs | 243 ++++++++- .../kernels/arc_kvquant/arc_kvquant.cu | 461 ++++++++++++++++++ mistralrs-quant/src/arc_kvquant/ffi.rs | 59 +++ mistralrs-quant/src/arc_kvquant/mod.rs | 380 +++++++++++++++ mistralrs-quant/src/lib.rs | 1 + 5 files changed, 1139 insertions(+), 5 deletions(-) create mode 100644 mistralrs-quant/kernels/arc_kvquant/arc_kvquant.cu create mode 100644 mistralrs-quant/src/arc_kvquant/ffi.rs create mode 100644 mistralrs-quant/src/arc_kvquant/mod.rs diff --git a/mistralrs-core/src/models/dsv4_kv_fp8.rs b/mistralrs-core/src/models/dsv4_kv_fp8.rs index 5752bebf5..773e95861 100644 --- a/mistralrs-core/src/models/dsv4_kv_fp8.rs +++ b/mistralrs-core/src/models/dsv4_kv_fp8.rs @@ -48,6 +48,7 @@ use candle_core::{DType, Device, Result, Tensor, D}; use float8::F8E4M3; +use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{OnceLock, RwLock}; /// Block width of the activation quantizer, matching the reference's @@ -61,13 +62,31 @@ const E4M3_MAX: f64 = 448.0; /// Which arithmetic produces the E4M3 code. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) enum KvQuantMode { + /// One fused device kernel per direction (`ArcQuant`, + /// `mistralrs_quant::arc_kvquant`). **The default.** Bit-identical to + /// [`Self::CpuExact`] — see that kernel's header for how each arithmetic + /// step is pinned to the candle op it replaces — but it is one launch + /// instead of ~11, and it removes the blocking device→host copy entirely. + /// + /// Falls back to [`Self::CpuExact`] when the tensor is not on CUDA, which + /// is what keeps the CPU unit tests below meaningful. + FusedDevice, /// Exact E4M3 via candle's CPU cast (candle has no CUDA `F8E4M3` cast — /// "named symbol not found"), at the price of one device sync per layer. + /// `ARC_KV_FP8_MODE=cpu`. CpuExact, /// On-device float arithmetic that reproduces E4M3's value grid with /// round-half-away-from-zero instead of round-half-to-even. Removes the /// sync; the FP8-trained model tolerates the sub-ULP-of-FP8 difference. /// `ARC_GPU_ACT_QUANT=1`. (RUN-161 throughput.) + /// + /// **Kept reachable, but it is NOT the fix and must not be shipped as one.** + /// Measured on H200 at B=1 it is *slower* than the sync it removes + /// (interleaved A1 66.99 / B1 67.71 / A2 68.05 / B2 70.30 ms/token, with + /// `kv_fp8_quant` 73.49 → 134.18 µs/call): it trades one blocking copy for + /// ~19 extra elementwise launches per layer, and this machine is bottlenecked + /// on op count, not on kernel speed. That measurement is the reason + /// [`Self::FusedDevice`] is one kernel rather than a chain of candle ops. GpuApprox, } @@ -79,15 +98,40 @@ impl KvQuantMode { pub(crate) fn from_env() -> Self { static MODE: OnceLock = OnceLock::new(); *MODE.get_or_init(|| { - if std::env::var_os("ARC_GPU_ACT_QUANT").is_some() { - Self::GpuApprox - } else { - Self::CpuExact + match std::env::var("ARC_KV_FP8_MODE") + .unwrap_or_default() + .to_ascii_lowercase() + .as_str() + { + "cpu" | "cpu_exact" => Self::CpuExact, + "gpu" | "gpu_approx" => Self::GpuApprox, + "fused" | "fused_device" => Self::FusedDevice, + _ if std::env::var_os("ARC_GPU_ACT_QUANT").is_some() => Self::GpuApprox, + _ => Self::FusedDevice, } }) } } +/// How many times the fused device kernels actually ran. +/// +/// D18/D33: a parity test that passes while these stay at zero has proved +/// nothing — it compared the candle path to itself. Every test that asserts the +/// fused path is correct also asserts these moved, and the fused path is the +/// only thing that touches them. +static FUSED_QUANT_CALLS: AtomicU64 = AtomicU64::new(0); +static FUSED_DEQUANT_CALLS: AtomicU64 = AtomicU64::new(0); + +/// Number of `quantize_k` calls served by the fused device kernel. +pub(crate) fn fused_quantize_calls() -> u64 { + FUSED_QUANT_CALLS.load(Ordering::Relaxed) +} + +/// Number of `V4PackedK::dequant` calls served by the fused device kernel. +pub(crate) fn fused_dequantize_calls() -> u64 { + FUSED_DEQUANT_CALLS.load(Ordering::Relaxed) +} + /// The 256 E4M3 values, indexed by code byte. Entry `i` is exactly what /// candle's `F8E4M3 -> F32` cast yields for the byte `i`, which is what makes /// [`V4PackedK::dequant`] bit-exact against the pre-packing implementation. @@ -165,7 +209,42 @@ impl V4PackedK { /// /// Bit-identical to what the pre-packing `act_quant_kv_nope` stored — see /// the module docs and `kv_fp8_roundtrip_is_bit_exact_vs_reference`. + /// + /// Served by the fused device kernel when it can be (default), and by the + /// candle op chain otherwise. The two produce the same bits, so the choice + /// is invisible to callers; it is exposed to tests through + /// [`Self::dequant_with`] so parity can be checked *between* them. pub(crate) fn dequant(&self, out_dtype: DType) -> Result { + self.dequant_with(KvQuantMode::from_env(), out_dtype) + } + + /// [`Self::dequant`] with the path chosen explicitly. + pub(crate) fn dequant_with(&self, mode: KvQuantMode, out_dtype: DType) -> Result { + if mode == KvQuantMode::FusedDevice + && out_dtype == self.side.dtype() + && mistralrs_quant::arc_kvquant::kv_fp8_fused_available(self.codes.device()) + { + let nope = self.codes.dim(D::Minus1)?; + let lut = e4m3_lut(self.codes.device())?; + let out = mistralrs_quant::arc_kvquant::kv_fp8_dequantize( + &self.codes, + &self.side, + &lut, + self.rope_dim, + KV_QUANT_BLOCK, + )?; + debug_assert_eq!(out.dim(D::Minus1)?, nope + self.rope_dim); + FUSED_DEQUANT_CALLS.fetch_add(1, Ordering::Relaxed); + return Ok(out); + } + self.dequant_candle(out_dtype) + } + + /// The candle op chain: ~13 launches, and the reference the fused kernel is + /// pinned against. Kept because it is the only implementation that runs on + /// CPU, and because a fused kernel with nothing to compare to is not a + /// verified kernel. + fn dequant_candle(&self, out_dtype: DType) -> Result { let (b, h, t, nope) = self.codes.dims4()?; let n_blocks = nope / KV_QUANT_BLOCK; let dev = self.codes.device(); @@ -222,6 +301,24 @@ pub(crate) fn quantize_k( } let n_blocks = nope / KV_QUANT_BLOCK; + // The whole point of this module's rewrite: one launch, no device→host + // round trip, and therefore nothing that blocks CUDA graph capture. Falls + // through to the candle chain on CPU (and on any dtype the kernel does not + // carry), which is what the parity tests compare against. + if mode == KvQuantMode::FusedDevice + && mistralrs_quant::arc_kvquant::kv_fp8_fused_available(k.device()) + && matches!(k.dtype(), DType::BF16 | DType::F16 | DType::F32) + { + let (codes, side) = + mistralrs_quant::arc_kvquant::kv_fp8_quantize(k, rope_dim, KV_QUANT_BLOCK)?; + FUSED_QUANT_CALLS.fetch_add(1, Ordering::Relaxed); + return Ok(Some(V4PackedK { + codes, + side, + rope_dim, + })); + } + let k_nope = k.narrow(D::Minus1, 0, nope)?; let k_rope = k.narrow(D::Minus1, nope, rope_dim)?; @@ -233,7 +330,10 @@ pub(crate) fn quantize_k( let scaled = kb.broadcast_div(&scale)?; let codes = match mode { - KvQuantMode::CpuExact => e4m3_codes_cpu(&scaled)?, + // `FusedDevice` only reaches here when the kernel could not serve the + // call (CPU tensor, or a dtype it does not carry). The exact path is + // the right fallback: it is what `FusedDevice` is bit-identical to. + KvQuantMode::FusedDevice | KvQuantMode::CpuExact => e4m3_codes_cpu(&scaled)?, KvQuantMode::GpuApprox => e4m3_codes_arith(&scaled)?, } .reshape((b, h, t, nope))?; @@ -490,6 +590,139 @@ mod tests { Ok(()) } + /// **The fused-kernel bit-parity proof, on hardware.** + /// + /// D14: this test is meaningless without a GPU, so it exits 2 (environment + /// failure) rather than passing when there is no CUDA device — a green run + /// with no device would be exactly the silent success this project has been + /// bitten by 13 times. + /// + /// D33: three independent quantities have to agree, and the mutant below + /// proves the comparison can fail. + /// + /// 1. `codes`/`side` bytes: fused kernel vs the candle chain + CPU cast. + /// 2. The reconstructed key tensor, computed **crosswise** — fused-quant + /// then candle-dequant against candle-quant then fused-dequant — so a + /// matching error would have to be present in both halves of two + /// different implementations. + /// 3. The full round trip against `reference_act_quant_kv_nope`, the + /// verbatim pre-packing implementation, run entirely on CPU. + /// + /// Plus the engagement assertion D18 demands: the fused call counters must + /// move, and must move only for the fused arm. A parity check on this exact + /// subsystem has already passed vacuously by comparing an implementation to + /// itself; only a launch counter caught it. + #[cfg(feature = "cuda")] + #[test] + fn kv_fp8_fused_is_bit_identical_to_cpu_exact() -> Result<()> { + const HEAD_DIM: usize = 512; + const ROPE: usize = 64; + + let dev = match Device::new_cuda(0) { + Ok(d) => d, + Err(e) => { + eprintln!("ENVFAIL: kv_fp8 fused parity needs a CUDA device: {e}"); + std::process::exit(2); + } + }; + + let k_cpu = fixture(2, 9, HEAD_DIM)?; + let k = k_cpu.to_device(&dev)?; + + let q0 = fused_quantize_calls(); + let packed_cpu = quantize_k(&k, ROPE, KvQuantMode::CpuExact)?.unwrap(); + assert_eq!( + fused_quantize_calls(), + q0, + "CpuExact must not reach the fused kernel — if it does, the two arms \ + are the same code and this test proves nothing" + ); + let packed_fused = quantize_k(&k, ROPE, KvQuantMode::FusedDevice)?.unwrap(); + assert_eq!( + fused_quantize_calls(), + q0 + 1, + "the fused arm did not reach the kernel: it silently fell back, so \ + any equality below is the candle path compared with itself" + ); + + // (1) The stored bytes. + let codes_cpu: Vec = packed_cpu.codes.flatten_all()?.to_vec1()?; + let codes_fused: Vec = packed_fused.codes.flatten_all()?.to_vec1()?; + assert_eq!( + codes_fused, codes_cpu, + "fused E4M3 codes differ from candle's CPU cast" + ); + assert_eq!( + bits(&packed_fused.side.to_device(&Device::Cpu)?)?, + bits(&packed_cpu.side.to_device(&Device::Cpu)?)?, + "fused `side` (rope tail + block amax) differs from the candle chain" + ); + + // (2) Crosswise reconstruction: each half of one implementation against + // the other half of the other. + let d0 = fused_dequantize_calls(); + let fused_then_candle = packed_fused.dequant_candle(DType::BF16)?; + assert_eq!( + fused_dequantize_calls(), + d0, + "dequant_candle must not reach the fused kernel" + ); + let candle_then_fused = packed_cpu.dequant_with(KvQuantMode::FusedDevice, DType::BF16)?; + assert_eq!( + fused_dequantize_calls(), + d0 + 1, + "the fused dequant silently fell back — nothing was tested" + ); + assert_eq!( + bits(&candle_then_fused.to_device(&Device::Cpu)?)?, + bits(&fused_then_candle.to_device(&Device::Cpu)?)?, + "fused dequant differs from the candle dequant chain" + ); + + // (3) Against the verbatim pre-packing implementation, on CPU. + let reference = reference_act_quant_kv_nope(&k_cpu, ROPE, false)?; + assert_eq!( + bits(&candle_then_fused.to_device(&Device::Cpu)?)?, + bits(&reference)?, + "the fused round trip does not reproduce what the pre-packing \ + implementation stored" + ); + + // D33: the comparison above must be able to fail. Same kernel, same + // scale, same layout, same reduction — only round-to-nearest-even + // replaced by truncation. + let (mutant_codes, mutant_side) = + mistralrs_quant::arc_kvquant::kv_fp8_quantize_mutant_for_test( + &k.contiguous()?, + ROPE, + KV_QUANT_BLOCK, + )?; + let mutant_codes: Vec = mutant_codes.flatten_all()?.to_vec1()?; + assert_eq!( + bits(&mutant_side.to_device(&Device::Cpu)?)?, + bits(&packed_cpu.side.to_device(&Device::Cpu)?)?, + "the mutant must differ ONLY in the rounding — if `side` differs too \ + it is not isolating the property under test" + ); + assert_ne!( + mutant_codes, codes_cpu, + "the negative control produced identical codes: this parity test \ + cannot fail and therefore proves nothing (D33)" + ); + let differing = mutant_codes + .iter() + .zip(codes_cpu.iter()) + .filter(|(a, b)| a != b) + .count(); + assert!( + differing * 20 > codes_cpu.len(), + "the negative control only perturbed {differing}/{} codes; a \ + comparison that weak could pass on luck", + codes_cpu.len() + ); + Ok(()) + } + /// The decode table must be the E4M3 grid itself, not an approximation of /// it: this is the single assumption `dequant`'s exactness rests on. #[test] diff --git a/mistralrs-quant/kernels/arc_kvquant/arc_kvquant.cu b/mistralrs-quant/kernels/arc_kvquant/arc_kvquant.cu new file mode 100644 index 000000000..7afb89f6e --- /dev/null +++ b/mistralrs-quant/kernels/arc_kvquant/arc_kvquant.cu @@ -0,0 +1,461 @@ +// Parent system: ArcQuant / TurboQuant — fused block-wise E4M3 activation +// quantizer for DeepSeek-V4's fused MQA key cache (ArcInfer / ArcKV / Fp8). +// +// WHY THIS FILE EXISTS +// ------------------- +// `dsv4_kv_fp8::e4m3_codes_cpu` quantized K by copying the scaled block to the +// HOST, casting there (candle has no CUDA F8E4M3 cast) and copying back. On the +// V4-Flash B=1 decode step that is 43 `cuMemcpyDtoHAsync_v2` per token at +// ~109 us of *host* time each (pageable staging => blocking) = 4.81 ms/token, +// and it makes CUDA graph capture impossible: a graph cannot record a blocking +// D2H. +// +// The obvious sync-free replacement (`KvQuantMode::GpuApprox`, ~19 extra +// elementwise candle ops per layer) was MEASURED SLOWER on H200 - in a machine +// bottlenecked on op count, one sync is cheaper than the launches it takes to +// avoid it. So the fix has to be *one* launch, not *fewer syncs*: these two +// kernels collapse the whole quantize (~11 candle ops) and dequantize (~13 +// candle ops) chains into a single kernel each. +// +// BIT-PARITY IS THE BAR - HOW IT IS OBTAINED +// ----------------------------------------- +// The stored bytes must be identical to what the CPU path already wrote, or +// every cache written before this change decodes differently. Each arithmetic +// step below therefore names the candle op it reproduces: +// +// amax = max |x| over the block <- abs() + max_keepdim() +// scale = amax * (float)(1/448) + 1e-12f <- `(amax / 448.0)?.affine(1,1e-12)` +// NOTE candle's `Tensor / f64` +// is `affine(1.0/rhs, 0.0)`, a +// MULTIPLY by the reciprocal, +// not a divide. +// code = E4M3(x / scale) <- broadcast_div + CPU cast +// value = lut[code] * scale <- index_select + broadcast_mul +// +// Two hazards are handled explicitly: +// +// 1. THIS CRATE COMPILES WITH `--use_fast_math` (mistralrs-quant/build.rs), so +// a bare `a / b` here would be `div.approx.f32` while candle-kernels (no +// fast math) emits IEEE `div.rn.f32`, and `-ftz=true` would flush denormals +// candle keeps. Every float operation below is therefore an explicit +// `__f*_rn` intrinsic, which the CUDA Math API defines as IEEE-754 and +// unaffected by `-prec-div`/`-ftz`/`-fmad`. `fabsf`/`fmaxf` are avoided in +// favour of integer ops on the bit pattern for the same reason. +// 2. The E4M3 rounding is a transcription of NVIDIA's +// `__nv_cvt_double_to_fp8(x, __NV_SATFINITE, __NV_E4M3)` (cuda_fp8.hpp), +// which is *also* what the Rust `float8` crate ports in `convert_to_fp8` +// and therefore what `F8E4M3::from_f32` - candle's CPU cast - computes. It +// is transcribed rather than #included so the parity claim is auditable +// against the Rust source in one place, and so it stays pure integer +// arithmetic that `--use_fast_math` cannot reach. +// +// The DEQUANT side does not convert code->float at all: it indexes the same +// 256-entry `f32` table the Rust side already builds with +// `F8E4M3::from_bits(i).to_f32()` and caches per device, so that half of the +// round trip is bit-exact by construction rather than by argument. + +#include +#include +#include + +namespace arc { + +// E4M3 max magnitude; `dsv4_kv_fp8::E4M3_MAX`. +#define ARC_KV_E4M3_MAX 448.0 +// The warp width the block reduction assumes. +#define ARC_KV_WARP 32 + +// --------------------------------------------------------------------------- +// Activation dtype <-> f32, matching candle's cast kernels exactly. +// f32 -> bf16 : candle `cast_f32_bf16` is `out[i] = inp[i]`, i.e. the +// `__nv_bfloat16(float)` constructor == `__float2bfloat16` +// (inline `cvt.rn.bf16.f32`, immune to -ftz). +// f32 -> f16 : likewise `__float2half` (inline `cvt.rn.f16.f32`). +// --------------------------------------------------------------------------- +template __device__ __forceinline__ float arc_to_f32(T v); +template <> __device__ __forceinline__ float arc_to_f32<__nv_bfloat16>(__nv_bfloat16 v) { + return __bfloat162float(v); +} +template <> __device__ __forceinline__ float arc_to_f32<__half>(__half v) { + return __half2float(v); +} +template <> __device__ __forceinline__ float arc_to_f32(float v) { return v; } + +template __device__ __forceinline__ T arc_from_f32(float v); +template <> __device__ __forceinline__ __nv_bfloat16 arc_from_f32<__nv_bfloat16>(float v) { + return __float2bfloat16(v); +} +template <> __device__ __forceinline__ __half arc_from_f32<__half>(float v) { + return __float2half(v); +} +template <> __device__ __forceinline__ float arc_from_f32(float v) { return v; } + +// --------------------------------------------------------------------------- +// scale = (amax / 448.0) then +1e-12, reproducing `dsv4_kv_fp8::block_scale`: +// (amax / E4M3_MAX)?.affine(1.0, 1e-12) +// candle lowers `Tensor / f64` to `affine(1.0/448.0, 0.0)` (bin_trait!(Div, .., +// |v| 1./v, ..)), i.e. a multiply by the f32-rounded reciprocal, and lowers +// `affine(mul, add)` to `x * mul + add` with `mul`/`add` first narrowed to the +// tensor dtype. `x * 1.0f` is exact, so the second affine is exactly an add, +// and FMA contraction of `x * mul + 0.0f` is exactly the multiply. Both are +// therefore reproduced by one `__fmul_rn` and one `__fadd_rn`. +// --------------------------------------------------------------------------- +__device__ __forceinline__ float arc_kv_block_scale(float amax) { + const float inv_max = (float)(1.0 / ARC_KV_E4M3_MAX); + const float eps = (float)1e-12; + return __fadd_rn(__fmul_rn(amax, inv_max), eps); +} + +// --------------------------------------------------------------------------- +// f32 -> E4M3 code byte. +// +// Transcription of NVIDIA `__nv_cvt_double_to_fp8(x, __NV_SATFINITE, +// __NV_E4M3)` (cuda_fp8.hpp), which the Rust `float8` crate ports verbatim as +// `convert_to_fp8`; `F8E4M3::from_f32(x)` is `convert_to_fp8(x as f64, +// SatFinite, E4M3)`, and `f32 -> f64` widening is exact. Every operation here +// is integer, so `--use_fast_math` cannot perturb it. +// +// Kept deliberately branch-shaped like the reference: the cost that mattered +// was a 109 us blocking memcpy, not a dozen integer ops. +// --------------------------------------------------------------------------- +__device__ __forceinline__ uint8_t arc_f32_to_e4m3_code(float xf) { + const double x = (double)xf; // exact + const uint64_t xbits = (uint64_t)__double_as_longlong(x); + + const uint8_t FP8_MAXNORM = 0x7EU; + const uint8_t FP8_MANTISSA_MASK = 0x07U; + const int FP8_EXP_BIAS = 7; + const int FP8_SIGNIFICAND_BITS = 4; + const uint64_t FP8_MINDENORM_O2 = 0x3F50000000000000ULL; + const uint64_t FP8_OVERFLOW_THRESHOLD = 0x407D000000000000ULL; + const uint64_t FP8_MINNORM = 0x3F90000000000000ULL; + const uint64_t DP_INF_BITS = 0x7FF0000000000000ULL; + const uint64_t FP8_DP_HALF_ULP = (uint64_t)1 << (53 - FP8_SIGNIFICAND_BITS - 1); + + const uint8_t sign = (uint8_t)((xbits >> 63) << 7); + const uint8_t exp = + (uint8_t)((int)((xbits >> 52) & 0x7FFULL) - 1023 + FP8_EXP_BIAS); + const uint8_t mantissa = + (uint8_t)((xbits >> (53 - FP8_SIGNIFICAND_BITS)) & (uint64_t)FP8_MANTISSA_MASK); + const uint64_t absx = xbits & 0x7FFFFFFFFFFFFFFFULL; + + uint8_t res; + if (absx <= FP8_MINDENORM_O2) { + // Zero or underflow. + res = 0U; + } else if (absx > DP_INF_BITS) { + // Preserve NaNs (E4M3 has a single NaN encoding). + res = 0x7FU; + } else if (absx > FP8_OVERFLOW_THRESHOLD) { + // SatFinite. + res = FP8_MAXNORM; + } else if (absx >= FP8_MINNORM) { + // Normal range, round-to-nearest-even. + res = (uint8_t)((uint8_t)(exp << (FP8_SIGNIFICAND_BITS - 1)) | mantissa); + const uint64_t round = xbits & ((FP8_DP_HALF_ULP << 1) - 1); + if ((round > FP8_DP_HALF_ULP) || + ((round == FP8_DP_HALF_ULP) && ((mantissa & 1U) != 0U))) { + res = (uint8_t)(res + 1U); + } + } else { + // Denormal range, round-to-nearest-even. + const uint8_t shift = (uint8_t)(1 - (int)exp); + const uint8_t man = (uint8_t)(mantissa | (uint8_t)(1U << (FP8_SIGNIFICAND_BITS - 1))); + res = (uint8_t)(man >> shift); + const uint64_t round = + (xbits | ((uint64_t)1 << (53 - 1))) & + ((FP8_DP_HALF_ULP << ((uint64_t)shift + 1)) - 1); + if ((round > (FP8_DP_HALF_ULP << (uint64_t)shift)) || + ((round == (FP8_DP_HALF_ULP << (uint64_t)shift)) && ((res & 1U) != 0U))) { + res = (uint8_t)(res + 1U); + } + } + return (uint8_t)(res | sign); +} + +// --------------------------------------------------------------------------- +// FUSED QUANTIZE. One launch replaces narrow/to_dtype/reshape/abs/max_keepdim/ +// affine/affine/broadcast_div/cast/reshape/cat. +// +// One CUDA block per token (a token is one (b, h, t) triple); one warp per +// 64-wide quant block. The amax reduction runs on the ABSOLUTE-VALUE BIT +// PATTERN as an unsigned integer: for non-negative IEEE floats the unsigned +// integer order is the float order, so `max` over `bits & 0x7fffffff` is +// exactly `max |x|` with no rounding, no `fabsf`, and no exposure to -ftz. +// --------------------------------------------------------------------------- +template +__global__ void arc_kv_fp8_quantize_kernel( + const T *__restrict__ k, // [ntok, head_dim] contiguous + uint8_t *__restrict__ codes, // [ntok, nope] + T *__restrict__ side, // [ntok, rope_dim + n_blocks] + const int head_dim, const int nope, const int rope_dim, const int n_blocks, + const int block_w, const long ntok) { + const int side_w = rope_dim + n_blocks; + const int lane = (int)(threadIdx.x & (ARC_KV_WARP - 1)); + const int warp = (int)(threadIdx.x / ARC_KV_WARP); + const int n_warps = (int)(blockDim.x / ARC_KV_WARP); + + for (long tok = (long)blockIdx.x; tok < ntok; tok += (long)gridDim.x) { + const T *krow = k + tok * (long)head_dim; + uint8_t *crow = codes + tok * (long)nope; + T *srow = side + tok * (long)side_w; + + for (int blk = warp; blk < n_blocks; blk += n_warps) { + const int base = blk * block_w; + + unsigned amax_bits = 0U; + for (int e = lane; e < block_w; e += ARC_KV_WARP) { + const unsigned b = __float_as_uint(arc_to_f32(krow[base + e])) & 0x7FFFFFFFU; + amax_bits = amax_bits > b ? amax_bits : b; + } +#pragma unroll + for (int off = ARC_KV_WARP / 2; off > 0; off >>= 1) { + const unsigned o = __shfl_xor_sync(0xFFFFFFFFU, amax_bits, off); + amax_bits = amax_bits > o ? amax_bits : o; + } + const float amax = __uint_as_float(amax_bits); + const float scale = arc_kv_block_scale(amax); + + for (int e = lane; e < block_w; e += ARC_KV_WARP) { + const float v = arc_to_f32(krow[base + e]); + crow[base + e] = arc_f32_to_e4m3_code(__fdiv_rn(v, scale)); + } + // `amax` IS one of the block's own elements, so narrowing it back to the + // activation dtype is exact - that is what lets dequant rebuild the + // identical scale. + if (lane == 0) { + srow[rope_dim + blk] = arc_from_f32(amax); + } + } + + // The RoPE'd tail is stored verbatim (RoPE is applied before the + // quantizer), and it is the FIRST `rope_dim` lanes of `side` - + // `cat(&[k_rope, amax_stored])`. + for (int i = (int)threadIdx.x; i < rope_dim; i += (int)blockDim.x) { + srow[i] = krow[nope + i]; + } + } +} + +// --------------------------------------------------------------------------- +// FUSED DEQUANTIZE. One launch replaces flatten/to_dtype/index_select/reshape/ +// narrow/to_dtype/reshape/affine/affine/broadcast_mul/reshape/to_dtype/narrow/ +// to_dtype/cat/contiguous. +// +// `lut` is the 256-entry f32 table the Rust side builds from +// `F8E4M3::from_bits(i).to_f32()` and caches per device - the SAME tensor the +// candle path passed to `index_select`, so the code->value half is bit-exact by +// construction. +// --------------------------------------------------------------------------- +template +__global__ void arc_kv_fp8_dequantize_kernel( + const uint8_t *__restrict__ codes, // [ntok, nope] + const T *__restrict__ side, // [ntok, rope_dim + n_blocks] + const float *__restrict__ lut, // [256] + T *__restrict__ out, // [ntok, head_dim] contiguous + const int head_dim, const int nope, const int rope_dim, const int n_blocks, + const int block_w, const long ntok) { + const int side_w = rope_dim + n_blocks; + const int lane = (int)(threadIdx.x & (ARC_KV_WARP - 1)); + const int warp = (int)(threadIdx.x / ARC_KV_WARP); + const int n_warps = (int)(blockDim.x / ARC_KV_WARP); + + for (long tok = (long)blockIdx.x; tok < ntok; tok += (long)gridDim.x) { + const uint8_t *crow = codes + tok * (long)nope; + const T *srow = side + tok * (long)side_w; + T *orow = out + tok * (long)head_dim; + + for (int blk = warp; blk < n_blocks; blk += n_warps) { + const int base = blk * block_w; + const float amax = arc_to_f32(srow[rope_dim + blk]); + const float scale = arc_kv_block_scale(amax); + for (int e = lane; e < block_w; e += ARC_KV_WARP) { + orow[base + e] = arc_from_f32(__fmul_rn(lut[crow[base + e]], scale)); + } + } + + // `cat(&[k_nope, k_rope])`: the RoPE'd tail follows the nope dims. + for (int i = (int)threadIdx.x; i < rope_dim; i += (int)blockDim.x) { + orow[nope + i] = srow[i]; + } + } +} + +} // namespace arc + +// --------------------------------------------------------------------------- +// C ABI. `stream_ptr` is candle's stream so the launch is recorded into a CUDA +// graph captured on it (cf. kvwrite.cu; rotary.cu's hardcoded stream 0 is NOT +// capturable). Grid is capped so a long prefill loops instead of launching an +// unbounded grid. +// --------------------------------------------------------------------------- +#define ARC_KV_MAX_BLOCKS 65535 + +#define ARC_KV_LAUNCH_QUANT(T) \ + arc::arc_kv_fp8_quantize_kernel<<>>( \ + reinterpret_cast(k), codes, reinterpret_cast(side), \ + head_dim, nope, rope_dim, n_blocks, block_w, ntok) + +#define ARC_KV_LAUNCH_DEQUANT(T) \ + arc::arc_kv_fp8_dequantize_kernel<<>>( \ + codes, reinterpret_cast(side), lut, \ + reinterpret_cast(out), head_dim, nope, rope_dim, n_blocks, block_w, \ + ntok) + +extern "C" void arc_kv_fp8_quantize( + const void *k, // [ntok, head_dim] activation dtype + uint8_t *codes, // [ntok, nope] + void *side, // [ntok, rope_dim + n_blocks] + int32_t head_dim, int32_t nope, int32_t rope_dim, int32_t n_blocks, + int32_t block_w, int64_t ntok, + void *stream_ptr, // cudaStream_t; null => default stream + uint32_t dtype // 0 => f16, 1 => bf16, 2 => f32 +) { + if (ntok <= 0 || n_blocks <= 0) { + return; + } + // One warp per quant block, capped at 256 threads (8 warps). + int warps = n_blocks < 8 ? n_blocks : 8; + const int threads = warps * ARC_KV_WARP; + long blocks = ntok < (long)ARC_KV_MAX_BLOCKS ? ntok : (long)ARC_KV_MAX_BLOCKS; + dim3 grid((unsigned)blocks); + dim3 block((unsigned)threads); + const cudaStream_t stream = reinterpret_cast(stream_ptr); + + if (dtype == 0) { + ARC_KV_LAUNCH_QUANT(__half); + } else if (dtype == 1) { + ARC_KV_LAUNCH_QUANT(__nv_bfloat16); + } else if (dtype == 2) { + ARC_KV_LAUNCH_QUANT(float); + } +} + +extern "C" void arc_kv_fp8_dequantize( + const uint8_t *codes, // [ntok, nope] + const void *side, // [ntok, rope_dim + n_blocks] + const float *lut, // [256] + void *out, // [ntok, head_dim] + int32_t head_dim, int32_t nope, int32_t rope_dim, int32_t n_blocks, + int32_t block_w, int64_t ntok, void *stream_ptr, uint32_t dtype) { + if (ntok <= 0 || n_blocks <= 0) { + return; + } + int warps = n_blocks < 8 ? n_blocks : 8; + const int threads = warps * ARC_KV_WARP; + long blocks = ntok < (long)ARC_KV_MAX_BLOCKS ? ntok : (long)ARC_KV_MAX_BLOCKS; + dim3 grid((unsigned)blocks); + dim3 block((unsigned)threads); + const cudaStream_t stream = reinterpret_cast(stream_ptr); + + if (dtype == 0) { + ARC_KV_LAUNCH_DEQUANT(__half); + } else if (dtype == 1) { + ARC_KV_LAUNCH_DEQUANT(__nv_bfloat16); + } else if (dtype == 2) { + ARC_KV_LAUNCH_DEQUANT(float); + } +} + +// --------------------------------------------------------------------------- +// D33 NEGATIVE CONTROL. Same as `arc_kv_fp8_quantize` except the rounding is +// truncated (round-toward-zero) instead of round-to-nearest-even. It exists so +// the bit-parity test can be shown to FAIL on a wrong kernel: a parity check on +// this exact subsystem has already passed vacuously by comparing an +// implementation to itself. +// +// Nothing in the serving path calls this - it is reachable only from +// `arc_kvquant::mutant_quantize_for_parity_test`. +// --------------------------------------------------------------------------- +namespace arc { + +// Identical to `arc_f32_to_e4m3_code` except that in the normal range it +// TRUNCATES the mantissa instead of rounding to nearest even - the single most +// plausible way to get an E4M3 quantizer subtly wrong. Everything else (the +// scale, the amax reduction, the layout, the rope copy) is untouched, so a +// parity test that survives this mutant is testing nothing. +__device__ __forceinline__ uint8_t arc_f32_to_e4m3_code_truncating(float xf) { + const double x = (double)xf; + const uint64_t xbits = (uint64_t)__double_as_longlong(x); + const uint64_t absx = xbits & 0x7FFFFFFFFFFFFFFFULL; + const uint64_t FP8_MINNORM = 0x3F90000000000000ULL; + const uint64_t FP8_OVERFLOW_THRESHOLD = 0x407D000000000000ULL; + if (absx < FP8_MINNORM || absx > FP8_OVERFLOW_THRESHOLD) { + return arc_f32_to_e4m3_code(xf); + } + const uint8_t exp = (uint8_t)((int)((xbits >> 52) & 0x7FFULL) - 1023 + 7); + const uint8_t mantissa = (uint8_t)((xbits >> 49) & 0x7ULL); + const uint8_t sign = (uint8_t)((xbits >> 63) << 7); + return (uint8_t)((uint8_t)((uint8_t)(exp << 3) | mantissa) | sign); +} + +template +__global__ void arc_kv_fp8_quantize_mutant_kernel( + const T *__restrict__ k, uint8_t *__restrict__ codes, T *__restrict__ side, + const int head_dim, const int nope, const int rope_dim, const int n_blocks, + const int block_w, const long ntok) { + const int side_w = rope_dim + n_blocks; + const int lane = (int)(threadIdx.x & (ARC_KV_WARP - 1)); + const int warp = (int)(threadIdx.x / ARC_KV_WARP); + const int n_warps = (int)(blockDim.x / ARC_KV_WARP); + + for (long tok = (long)blockIdx.x; tok < ntok; tok += (long)gridDim.x) { + const T *krow = k + tok * (long)head_dim; + uint8_t *crow = codes + tok * (long)nope; + T *srow = side + tok * (long)side_w; + for (int blk = warp; blk < n_blocks; blk += n_warps) { + const int base = blk * block_w; + unsigned amax_bits = 0U; + for (int e = lane; e < block_w; e += ARC_KV_WARP) { + const unsigned b = __float_as_uint(arc_to_f32(krow[base + e])) & 0x7FFFFFFFU; + amax_bits = amax_bits > b ? amax_bits : b; + } +#pragma unroll + for (int off = ARC_KV_WARP / 2; off > 0; off >>= 1) { + const unsigned o = __shfl_xor_sync(0xFFFFFFFFU, amax_bits, off); + amax_bits = amax_bits > o ? amax_bits : o; + } + const float amax = __uint_as_float(amax_bits); + const float scale = arc_kv_block_scale(amax); + for (int e = lane; e < block_w; e += ARC_KV_WARP) { + const float v = arc_to_f32(krow[base + e]); + crow[base + e] = arc_f32_to_e4m3_code_truncating(__fdiv_rn(v, scale)); + } + if (lane == 0) { + srow[rope_dim + blk] = arc_from_f32(amax); + } + } + for (int i = (int)threadIdx.x; i < rope_dim; i += (int)blockDim.x) { + srow[i] = krow[nope + i]; + } + } +} + +} // namespace arc + +#define ARC_KV_LAUNCH_MUTANT(T) \ + arc::arc_kv_fp8_quantize_mutant_kernel<<>>( \ + reinterpret_cast(k), codes, reinterpret_cast(side), \ + head_dim, nope, rope_dim, n_blocks, block_w, ntok) + +extern "C" void arc_kv_fp8_quantize_mutant( + const void *k, uint8_t *codes, void *side, int32_t head_dim, int32_t nope, + int32_t rope_dim, int32_t n_blocks, int32_t block_w, int64_t ntok, + void *stream_ptr, uint32_t dtype) { + if (ntok <= 0 || n_blocks <= 0) { + return; + } + int warps = n_blocks < 8 ? n_blocks : 8; + const int threads = warps * ARC_KV_WARP; + long blocks = ntok < (long)ARC_KV_MAX_BLOCKS ? ntok : (long)ARC_KV_MAX_BLOCKS; + dim3 grid((unsigned)blocks); + dim3 block((unsigned)threads); + const cudaStream_t stream = reinterpret_cast(stream_ptr); + if (dtype == 0) { + ARC_KV_LAUNCH_MUTANT(__half); + } else if (dtype == 1) { + ARC_KV_LAUNCH_MUTANT(__nv_bfloat16); + } else if (dtype == 2) { + ARC_KV_LAUNCH_MUTANT(float); + } +} diff --git a/mistralrs-quant/src/arc_kvquant/ffi.rs b/mistralrs-quant/src/arc_kvquant/ffi.rs new file mode 100644 index 000000000..dafc5bd42 --- /dev/null +++ b/mistralrs-quant/src/arc_kvquant/ffi.rs @@ -0,0 +1,59 @@ +use core::ffi::c_void; + +extern "C" { + /// Fused block-wise E4M3 quantize. See `kernels/arc_kvquant/arc_kvquant.cu`. + /// + /// * `k` - `[ntok, head_dim]` contiguous, activation dtype. + /// * `codes` - `[ntok, nope]` U8 out. + /// * `side` - `[ntok, rope_dim + n_blocks]` activation dtype out + /// (`cat(&[k_rope, amax])`). + /// * `stream` - candle's stream, so the launch is recordable into a CUDA + /// graph captured on it. + /// * `dtype` - 0 => f16, 1 => bf16, 2 => f32. + pub(crate) fn arc_kv_fp8_quantize( + k: *const c_void, + codes: *mut u8, + side: *mut c_void, + head_dim: i32, + nope: i32, + rope_dim: i32, + n_blocks: i32, + block_w: i32, + ntok: i64, + stream: *mut c_void, + dtype: u32, + ); + + /// Fused dequantize back to `[ntok, head_dim]` (`cat(&[k_nope, k_rope])`). + /// `lut` is the 256-entry F32 `F8E4M3::from_bits(i).to_f32()` table. + pub(crate) fn arc_kv_fp8_dequantize( + codes: *const u8, + side: *const c_void, + lut: *const f32, + out: *mut c_void, + head_dim: i32, + nope: i32, + rope_dim: i32, + n_blocks: i32, + block_w: i32, + ntok: i64, + stream: *mut c_void, + dtype: u32, + ); + + /// D33 negative control: `arc_kv_fp8_quantize` with round-to-nearest-even + /// replaced by truncation. Test-only; nothing in the serving path calls it. + pub(crate) fn arc_kv_fp8_quantize_mutant( + k: *const c_void, + codes: *mut u8, + side: *mut c_void, + head_dim: i32, + nope: i32, + rope_dim: i32, + n_blocks: i32, + block_w: i32, + ntok: i64, + stream: *mut c_void, + dtype: u32, + ); +} diff --git a/mistralrs-quant/src/arc_kvquant/mod.rs b/mistralrs-quant/src/arc_kvquant/mod.rs new file mode 100644 index 000000000..887043b68 --- /dev/null +++ b/mistralrs-quant/src/arc_kvquant/mod.rs @@ -0,0 +1,380 @@ +//! Parent system: ArcQuant / TurboQuant — fused block-wise E4M3 activation +//! quantizer for the DeepSeek-V4 fused MQA key cache (ArcInfer / ArcKV / Fp8). +//! +//! One launch each for quantize and dequantize, replacing the ~11 and ~13 +//! candle ops the caller used to issue — and, critically, replacing the +//! device→host→device round trip that `dsv4_kv_fp8::e4m3_codes_cpu` performs +//! because candle has no CUDA `F8E4M3` cast. That round trip is 43 blocking +//! `cuMemcpyDtoHAsync_v2` per V4 decode token and it makes CUDA graph capture +//! impossible: a graph cannot record a blocking D2H. +//! +//! The byte format is unchanged — same `codes`/`side` layout, same bits. See +//! `kernels/arc_kvquant/arc_kvquant.cu` for how bit-parity with the CPU path is +//! obtained, and for the deliberate mutant used to prove the parity test can +//! fail. + +#[cfg(feature = "cuda")] +mod ffi; + +#[cfg(feature = "cuda")] +mod cuda_impl { + use candle_core::{ + cuda_backend::{cudarc::driver::DeviceRepr, CudaDType}, + CudaDevice, CudaStorage, DType, Device, Result, Shape, Storage, Tensor, + }; + use core::ffi::c_void; + use half::{bf16, f16}; + + use crate::utils::slice_ptr; + + /// `dtype` discriminant shared with the C ABI. + fn dtype_id(dt: DType) -> Result { + match dt { + DType::F16 => Ok(0), + DType::BF16 => Ok(1), + DType::F32 => Ok(2), + other => candle_core::bail!("arc-kvquant: unsupported activation dtype {other:?}"), + } + } + + /// Geometry shared by both kernels, validated once. + struct Geom { + b: usize, + h: usize, + t: usize, + head_dim: usize, + nope: usize, + rope_dim: usize, + n_blocks: usize, + block_w: usize, + } + + impl Geom { + fn ntok(&self) -> usize { + self.b * self.h * self.t + } + fn side_w(&self) -> usize { + self.rope_dim + self.n_blocks + } + } + + fn cuda_device(t: &Tensor, what: &str) -> Result { + match t.device() { + Device::Cuda(d) => Ok(d.clone()), + _ => candle_core::bail!("arc-kvquant: {what} must live on CUDA"), + } + } + + /// Whether the fused kernels can serve `dev`. Callers fall back to the + /// candle op chain when this is false rather than failing. + pub fn kv_fp8_fused_available(dev: &Device) -> bool { + matches!(dev, Device::Cuda(_)) + } + + fn geometry(k: &Tensor, rope_dim: usize, block_w: usize) -> Result { + let (b, h, t, head_dim) = k.dims4()?; + if rope_dim > head_dim { + candle_core::bail!("arc-kvquant: rope_dim {rope_dim} > head_dim {head_dim}"); + } + let nope = head_dim - rope_dim; + if block_w == 0 || nope == 0 || nope % block_w != 0 { + candle_core::bail!( + "arc-kvquant: nope {nope} is not a whole number of {block_w}-wide blocks" + ); + } + Ok(Geom { + b, + h, + t, + head_dim, + nope, + rope_dim, + n_blocks: nope / block_w, + block_w, + }) + } + + /// Fused block-wise E4M3 quantize of `k`'s non-RoPE dims. + /// + /// `k` is `[B, H, T, head_dim]` in the activation dtype. Returns + /// `(codes, side)`: + /// + /// * `codes` — `[B, H, T, head_dim - rope_dim]` U8, one E4M3 code per dim; + /// * `side` — `[B, H, T, rope_dim + n_blocks]` in `k`'s dtype: the RoPE'd + /// tail verbatim followed by each 64-wide block's `amax`. + /// + /// Bit-identical to `dsv4_kv_fp8::quantize_k` under `KvQuantMode::CpuExact`. + /// That is the contract; `kv_fp8_fused_is_bit_identical_to_cpu_exact` pins + /// it on hardware, and the mutant below proves that test can fail. + pub fn kv_fp8_quantize( + k: &Tensor, + rope_dim: usize, + block_w: usize, + ) -> Result<(Tensor, Tensor)> { + quantize_inner(k, rope_dim, block_w, false) + } + + /// D33 negative control: identical to [`kv_fp8_quantize`] except the E4M3 + /// rounding truncates instead of rounding to nearest even. Exists purely so + /// the parity test can be shown to fail. Nothing in the serving path calls + /// it. + pub fn kv_fp8_quantize_mutant_for_test( + k: &Tensor, + rope_dim: usize, + block_w: usize, + ) -> Result<(Tensor, Tensor)> { + quantize_inner(k, rope_dim, block_w, true) + } + + fn quantize_inner( + k: &Tensor, + rope_dim: usize, + block_w: usize, + mutant: bool, + ) -> Result<(Tensor, Tensor)> { + let g = geometry(k, rope_dim, block_w)?; + let dev = cuda_device(k, "k")?; + let did = dtype_id(k.dtype())?; + let k = k.contiguous()?; + match k.dtype() { + DType::F16 => quantize_t::(&k, &g, &dev, did, mutant), + DType::BF16 => quantize_t::(&k, &g, &dev, did, mutant), + DType::F32 => quantize_t::(&k, &g, &dev, did, mutant), + other => candle_core::bail!("arc-kvquant: unsupported activation dtype {other:?}"), + } + } + + fn quantize_t( + k: &Tensor, + g: &Geom, + dev: &CudaDevice, + did: u32, + mutant: bool, + ) -> Result<(Tensor, Tensor)> { + let ntok = g.ntok(); + // `alloc` rather than `alloc_zeros`: both outputs are fully written, and + // the memset would be one more device op in the exact place where op + // count is the disease. + let codes_buf = unsafe { dev.alloc::(ntok * g.nope)? }; + let side_buf = unsafe { dev.alloc::(ntok * g.side_w())? }; + + let (k_storage, k_layout) = k.storage_and_layout(); + let k_s = match &*k_storage { + Storage::Cuda(s) => s, + _ => candle_core::bail!("arc-kvquant: k must be CUDA storage"), + }; + let (k_ptr, _k_guard) = slice_ptr(k_s.as_cuda_slice::()?, k_layout.start_offset()); + let (codes_ptr, _codes_guard) = slice_ptr(&codes_buf, 0); + let (side_ptr, _side_guard) = slice_ptr(&side_buf, 0); + + let launch = if mutant { + super::ffi::arc_kv_fp8_quantize_mutant + } else { + super::ffi::arc_kv_fp8_quantize + }; + unsafe { + launch( + k_ptr as *const c_void, + codes_ptr as *mut u8, + side_ptr as *mut c_void, + g.head_dim as i32, + g.nope as i32, + g.rope_dim as i32, + g.n_blocks as i32, + g.block_w as i32, + ntok as i64, + dev.cuda_stream().cu_stream() as *mut c_void, + did, + ) + }; + + drop(_k_guard); + drop(_codes_guard); + drop(_side_guard); + drop(k_storage); + + let codes = Tensor::from(( + Storage::Cuda(CudaStorage::wrap_cuda_slice(codes_buf, dev.clone())), + Shape::from_dims(&[g.b, g.h, g.t, g.nope]), + )); + let side = Tensor::from(( + Storage::Cuda(CudaStorage::wrap_cuda_slice(side_buf, dev.clone())), + Shape::from_dims(&[g.b, g.h, g.t, g.side_w()]), + )); + Ok((codes, side)) + } + + /// Fused dequantize: `codes` + `side` back to `[B, H, T, head_dim]` in + /// `side`'s dtype. + /// + /// `lut` is the 256-entry F32 table built from + /// `F8E4M3::from_bits(i).to_f32()` — the *same* tensor the candle path fed + /// to `index_select`, which is what makes the code→value half of the round + /// trip bit-exact by construction rather than by argument. + pub fn kv_fp8_dequantize( + codes: &Tensor, + side: &Tensor, + lut: &Tensor, + rope_dim: usize, + block_w: usize, + ) -> Result { + let (b, h, t, nope) = codes.dims4()?; + let (sb, sh, st, side_w) = side.dims4()?; + if (sb, sh, st) != (b, h, t) { + candle_core::bail!( + "arc-kvquant: side dims {:?} do not match codes {:?}", + side.dims(), + codes.dims() + ); + } + if block_w == 0 || nope == 0 || nope % block_w != 0 { + candle_core::bail!("arc-kvquant: nope {nope} not a multiple of block_w {block_w}"); + } + let n_blocks = nope / block_w; + if side_w != rope_dim + n_blocks { + candle_core::bail!( + "arc-kvquant: side width {side_w} != rope_dim {rope_dim} + n_blocks {n_blocks}" + ); + } + if codes.dtype() != DType::U8 { + candle_core::bail!("arc-kvquant: codes must be U8, got {:?}", codes.dtype()); + } + if lut.dtype() != DType::F32 || lut.elem_count() != 256 { + candle_core::bail!("arc-kvquant: lut must be 256 F32 entries"); + } + let g = Geom { + b, + h, + t, + head_dim: nope + rope_dim, + nope, + rope_dim, + n_blocks, + block_w, + }; + let dev = cuda_device(side, "side")?; + let did = dtype_id(side.dtype())?; + // A `narrow` on the sequence dim leaves a contiguous view (with a start + // offset) whenever B*H == 1 — the decode case this exists for — so + // `contiguous()` is free there. For B*H > 1 it materialises the window, + // which is what threading 4-D strides into the kernel would avoid; that + // is a batch-path follow-up, not a correctness gap. + let codes = codes.contiguous()?; + let side = side.contiguous()?; + let lut = lut.contiguous()?; + match side.dtype() { + DType::F16 => dequantize_t::(&codes, &side, &lut, &g, &dev, did), + DType::BF16 => dequantize_t::(&codes, &side, &lut, &g, &dev, did), + DType::F32 => dequantize_t::(&codes, &side, &lut, &g, &dev, did), + other => candle_core::bail!("arc-kvquant: unsupported activation dtype {other:?}"), + } + } + + fn dequantize_t( + codes: &Tensor, + side: &Tensor, + lut: &Tensor, + g: &Geom, + dev: &CudaDevice, + did: u32, + ) -> Result { + let ntok = g.ntok(); + let out_buf = unsafe { dev.alloc::(ntok * g.head_dim)? }; + + let (codes_storage, codes_layout) = codes.storage_and_layout(); + let codes_s = match &*codes_storage { + Storage::Cuda(s) => s, + _ => candle_core::bail!("arc-kvquant: codes must be CUDA storage"), + }; + let (side_storage, side_layout) = side.storage_and_layout(); + let side_s = match &*side_storage { + Storage::Cuda(s) => s, + _ => candle_core::bail!("arc-kvquant: side must be CUDA storage"), + }; + let (lut_storage, lut_layout) = lut.storage_and_layout(); + let lut_s = match &*lut_storage { + Storage::Cuda(s) => s, + _ => candle_core::bail!("arc-kvquant: lut must be CUDA storage"), + }; + + let (codes_ptr, _codes_guard) = + slice_ptr(codes_s.as_cuda_slice::()?, codes_layout.start_offset()); + let (side_ptr, _side_guard) = + slice_ptr(side_s.as_cuda_slice::()?, side_layout.start_offset()); + let (lut_ptr, _lut_guard) = + slice_ptr(lut_s.as_cuda_slice::()?, lut_layout.start_offset()); + let (out_ptr, _out_guard) = slice_ptr(&out_buf, 0); + + unsafe { + super::ffi::arc_kv_fp8_dequantize( + codes_ptr as *const u8, + side_ptr as *const c_void, + lut_ptr as *const f32, + out_ptr as *mut c_void, + g.head_dim as i32, + g.nope as i32, + g.rope_dim as i32, + g.n_blocks as i32, + g.block_w as i32, + ntok as i64, + dev.cuda_stream().cu_stream() as *mut c_void, + did, + ) + }; + + drop(_codes_guard); + drop(_side_guard); + drop(_lut_guard); + drop(_out_guard); + drop(codes_storage); + drop(side_storage); + drop(lut_storage); + + Ok(Tensor::from(( + Storage::Cuda(CudaStorage::wrap_cuda_slice(out_buf, dev.clone())), + Shape::from_dims(&[g.b, g.h, g.t, g.head_dim]), + ))) + } +} + +#[cfg(feature = "cuda")] +pub use cuda_impl::*; + +#[cfg(not(feature = "cuda"))] +mod stub { + use candle_core::{Device, Result, Tensor}; + + pub fn kv_fp8_quantize( + _k: &Tensor, + _rope_dim: usize, + _block_w: usize, + ) -> Result<(Tensor, Tensor)> { + candle_core::bail!("arc-kvquant: fused FP8 KV quantize requires the `cuda` feature") + } + + pub fn kv_fp8_quantize_mutant_for_test( + _k: &Tensor, + _rope_dim: usize, + _block_w: usize, + ) -> Result<(Tensor, Tensor)> { + candle_core::bail!("arc-kvquant: fused FP8 KV quantize requires the `cuda` feature") + } + + pub fn kv_fp8_dequantize( + _codes: &Tensor, + _side: &Tensor, + _lut: &Tensor, + _rope_dim: usize, + _block_w: usize, + ) -> Result { + candle_core::bail!("arc-kvquant: fused FP8 KV dequantize requires the `cuda` feature") + } + + /// Never available without the `cuda` feature. + pub fn kv_fp8_fused_available(_dev: &Device) -> bool { + false + } +} + +#[cfg(not(feature = "cuda"))] +pub use stub::*; diff --git a/mistralrs-quant/src/lib.rs b/mistralrs-quant/src/lib.rs index bbe893b15..2bb42c3bd 100644 --- a/mistralrs-quant/src/lib.rs +++ b/mistralrs-quant/src/lib.rs @@ -16,6 +16,7 @@ use pertensor_fp8::pertensor_fp8_linear_b; mod metal_kernels; mod afq; +pub mod arc_kvquant; mod bitsandbytes; mod blockwise_fp8; pub mod calibration; From e10cdea8c149b4fd20a51ab0d9aea903a3c43533 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Tue, 18 Aug 2026 10:40:10 +0100 Subject: [PATCH 02/27] fix(ArcKV/Fp8): inline-PTX IEEE ops (fast-math rewrote __f*_rn to .ftz) + box paths The emitted PTX showed nvcc 13.1 turning __fmul_rn/__fadd_rn/__fdiv_rn into mul.rn.ftz.f32 / div.rn.ftz.f32 / add.rn.ftz.f32 under --use_fast_math, which candle-kernels (no fast math) does not do. Replaced with inline PTX, which no optimisation flag rewrites, so 'grep -c .ftz.f32' over the PTX is the audit. Co-Authored-By: Claude Opus 5 (1M context) --- arc-tools/measure_kv_fp8_fused.sh | 223 ++++++++++++++++++ .../kernels/arc_kvquant/arc_kvquant.cu | 55 ++++- 2 files changed, 268 insertions(+), 10 deletions(-) create mode 100755 arc-tools/measure_kv_fp8_fused.sh diff --git a/arc-tools/measure_kv_fp8_fused.sh b/arc-tools/measure_kv_fp8_fused.sh new file mode 100755 index 000000000..fcf03815a --- /dev/null +++ b/arc-tools/measure_kv_fp8_fused.sh @@ -0,0 +1,223 @@ +#!/usr/bin/env bash +# ArcKV/Fp8 — measure the fused E4M3 quantize+dequantize kernel against the +# CPU round trip it replaces. Runs ON THE BOX. Nothing here is a proxy: every +# number is produced by the real binary on the real model. +# +# WHAT IT PRODUCES (the four numbers the change is judged on) +# 1. `cuMemcpyDtoHAsync_v2` calls per decode step, before and after. +# The 44/step are the thing being removed and the reason CUDA graph +# capture is impossible. NOTE `*Synchronize*` is 0.0 calls/step in this +# workload — counting `cudaStreamSynchronize` reports "no syncs" and is +# wrong. Count the D2H. +# 2. Kernel launches per decode step, before and after. Op count is the +# disease (9,131/token, median kernel 1.18 us), so this matters more than +# wall time. +# 3. Interleaved A-B-A-B ms/token with the monotonic drift stated. This box +# showed 3.3% drift across one run — larger than a real arm difference — +# so A-then-B is a fabricated comparison. +# 4. Bit-parity: the GPU-gated `kv_fp8_fused_is_bit_identical_to_cpu_exact` +# test, plus its negative control. +# +# RULES ENFORCED HERE, NOT ASSUMED +# * The box is shared with the ArcGraph chain. Every timing leg holds +# /root/.arc-bench.lock under flock; a contended number is a fabricated +# number. +# * Exclusivity is asserted BEFORE AND AFTER every leg from +# `nvidia-smi --query-compute-apps` (the lock file is not an occupancy +# signal — an 87 GB server has been resident with no lock at all, and a +# 77 GB server has run with the lock reading FREE). A neighbour appearing +# mid-leg aborts the run. +# * Environment failure exits 2. A failed measurement exits 1. They are not +# the same thing and must never be reported as the same thing. +# * Engagement is asserted before any null/neutral result is believed: if +# the arm that is supposed to use the fused kernel shows the same D2H +# count as the arm that is not, the run aborts rather than reporting "no +# difference". +set -uo pipefail + +REPO="${REPO:-/root/arc-wt/fp8}" +MODEL="${MODEL:-deepseek-ai/DeepSeek-V4-Flash}" +ARCH="${ARCH:-deepseekv4}" +ISQ="${ISQ:-qtip2}" +PROMPT_LEN="${PROMPT_LEN:-128}" +GEN_LEN="${GEN_LEN:-64}" +MAX_SEQ_LEN="${MAX_SEQ_LEN:-1024}" +OUT="${OUT:-/root/kvfp8-fused}" +LOCK="${LOCK:-/root/locks/bench.lock}" +BIN="${BIN:-${CARGO_TARGET_DIR:-$REPO/target}/release/mistralrs}" +REPS="${REPS:-2}" # A-B-A-B => REPS=2 + +mkdir -p "$OUT" + +envfail() { + echo "ENVFAIL: $*" >&2 + exit 2 +} +fail() { + echo "FAIL: $*" >&2 + exit 1 +} + +command -v nvidia-smi >/dev/null 2>&1 || envfail "no nvidia-smi" +command -v nsys >/dev/null 2>&1 || echo "WARN: no nsys; the D2H/launch counts will be skipped" >&2 +[ -x "$BIN" ] || envfail "no mistralrs binary at $BIN (build with --features 'cuda flash-attn')" + +# --------------------------------------------------------------------------- +# Exclusivity. Bracket every leg: a V4 load shows near-zero VRAM for most of a +# minute, so "compute-apps empty" sampled during a neighbour's load is +# indistinguishable from an idle box. +# --------------------------------------------------------------------------- +MYPID="" +assert_exclusive() { + local where="$1" + local apps + apps=$(nvidia-smi --query-compute-apps=pid --format=csv,noheader | tr -d ' ' | sort -u | tr '\n' ',') + apps="${apps%,}" + if [ -n "$MYPID" ]; then + [ "$apps" = "$MYPID" ] || envfail "[$where] compute-apps=[$apps], expected only $MYPID" + else + [ -z "$apps" ] || envfail "[$where] box not idle: compute-apps=[$apps]" + fi +} + +# --------------------------------------------------------------------------- +# One measured leg. $1 = arm label, $2 = ARC_KV_FP8_MODE value ("" => default, +# which is the fused kernel), $3 = log path, $4 = "nsys" to trace. +# --------------------------------------------------------------------------- +run_leg() { + local arm="$1" mode="$2" log="$3" trace="${4:-}" + local -a cmd=("$BIN" bench -m "$MODEL" -a "$ARCH" --isq "$ISQ" + --prompt-len "$PROMPT_LEN" --gen-len "$GEN_LEN" --max-seq-len "$MAX_SEQ_LEN") + + assert_exclusive "pre-$arm" + ( + if [ -n "$mode" ]; then export ARC_KV_FP8_MODE="$mode"; else unset ARC_KV_FP8_MODE; fi + unset ARC_GPU_ACT_QUANT + if [ "$trace" = "nsys" ]; then + nsys profile -t cuda -o "$OUT/$arm" --force-overwrite true \ + --cuda-memory-usage false "${cmd[@]}" + else + "${cmd[@]}" + fi + ) >"$log" 2>&1 & + local pid=$! + # `pgrep -f`/`pkill -f` would match this script's own command line; use the + # job pid we already hold instead. + MYPID="" + wait "$pid" + local rc=$? + MYPID="" + assert_exclusive "post-$arm" + [ $rc -eq 0 ] || fail "$arm leg exited $rc; see $log" +} + +# --------------------------------------------------------------------------- +# ms/token out of a bench log. Fails loudly rather than emitting an empty +# string that would silently become a "0.0 ms/token improvement". +# --------------------------------------------------------------------------- +decode_ms_per_token() { + local log="$1" + local tps + tps=$(grep -oE 'tok_per_s_decode[^0-9]*[0-9]+\.[0-9]+' "$log" | tail -1 | + grep -oE '[0-9]+\.[0-9]+$') + [ -n "$tps" ] || tps=$(grep -oiE 'decode[^0-9]*([0-9]+\.[0-9]+) *tok' "$log" | tail -1 | + grep -oE '[0-9]+\.[0-9]+') + [ -n "$tps" ] || fail "no decode throughput in $log — do not report a delta from a missing number" + awk -v t="$tps" 'BEGIN { printf "%.3f", 1000.0 / t }' +} + +# --------------------------------------------------------------------------- +# 0. Bit-parity + its negative control. Cheap, and it gates everything else: +# a faster kernel that stores different bytes is not a faster kernel. +# --------------------------------------------------------------------------- +echo "=== 0. bit-parity (GPU-gated test + D33 negative control) ===" +( + cd "$REPO" || envfail "no repo at $REPO" + cargo test -p mistralrs-core --features "cuda flash-attn" --lib \ + kv_fp8_fused_is_bit_identical_to_cpu_exact -- --nocapture --exact \ + models::dsv4_kv_fp8::tests::kv_fp8_fused_is_bit_identical_to_cpu_exact +) 2>&1 | tee "$OUT/parity.log" +prc=${PIPESTATUS[0]} +# `set -e` does not survive a pipe; the status is read explicitly. +[ "$prc" -eq 2 ] && envfail "parity test could not find a CUDA device" +[ "$prc" -eq 0 ] || fail "bit-parity FAILED — the fused kernel does not store what the CPU path stores" +grep -q "test result: ok. 1 passed" "$OUT/parity.log" || + fail "parity test reported no results — 'no failures' is not 'ran'" +echo "bit-parity: PASS" + +# --------------------------------------------------------------------------- +# 1 + 2. D2H count and launch count per decode step, both arms, under nsys. +# --------------------------------------------------------------------------- +if command -v nsys >/dev/null 2>&1; then + echo "=== 1+2. nsys: DtoH copies and kernel launches per step ===" + flock "$LOCK" bash -c "$(declare -f run_leg assert_exclusive envfail fail); \ + OUT='$OUT' BIN='$BIN' MODEL='$MODEL' ARCH='$ARCH' ISQ='$ISQ' \ + PROMPT_LEN='$PROMPT_LEN' GEN_LEN='$GEN_LEN' MAX_SEQ_LEN='$MAX_SEQ_LEN' \ + bash -c 'true'" || true + + for arm in before after; do + mode=""; [ "$arm" = "before" ] && mode="cpu" + ( + flock 9 || envfail "could not take the box lock $LOCK" + run_leg "$arm" "$mode" "$OUT/$arm.nsys.log" nsys + ) 9>"$LOCK" + nsys stats --report cuda_api_sum --format csv "$OUT/$arm.nsys-rep" \ + >"$OUT/$arm.api.csv" 2>"$OUT/$arm.api.err" || + echo "WARN: nsys stats cuda_api_sum failed for $arm" >&2 + nsys stats --report cuda_gpu_kern_sum --format csv "$OUT/$arm.nsys-rep" \ + >"$OUT/$arm.kern.csv" 2>>"$OUT/$arm.api.err" || + echo "WARN: nsys stats cuda_gpu_kern_sum failed for $arm" >&2 + + d2h=$(awk -F, '/cuMemcpyDtoHAsync_v2/ { gsub(/"/,"",$3); s+=$3 } END { print s+0 }' "$OUT/$arm.api.csv") + launches=$(awk -F, 'NR>1 { gsub(/"/,"",$3); s+=$3 } END { print s+0 }' "$OUT/$arm.kern.csv") + echo "$arm: cuMemcpyDtoHAsync_v2=$d2h kernel_launches=$launches (steps=$GEN_LEN)" + echo "$arm $d2h $launches" >>"$OUT/counts.txt" + done + + # Engagement (D18): the two arms MUST differ in D2H count. Equal counts mean + # the fused arm never engaged, and every later number is the same code twice. + b=$(awk '$1=="before"{print $2}' "$OUT/counts.txt") + a=$(awk '$1=="after"{print $2}' "$OUT/counts.txt") + [ -n "$b" ] && [ -n "$a" ] || fail "missing D2H counts — no results is not a null result" + [ "$a" -lt "$b" ] || fail "ENGAGEMENT: after-arm D2H ($a) is not below before-arm ($b); the fused kernel did not run" + awk -v b="$b" -v a="$a" -v n="$GEN_LEN" 'BEGIN { + printf "D2H per step: before %.2f -> after %.2f (removed %.2f/step)\n", b/n, a/n, (b-a)/n }' +fi + +# --------------------------------------------------------------------------- +# 3. Interleaved A-B-A-B ms/token. A = CPU round trip, B = fused kernel. +# The drift is reported alongside the delta, because on this box a 3.3% +# monotonic drift once exceeded the arm difference. +# --------------------------------------------------------------------------- +echo "=== 3. interleaved A-B-A-B ms/token ===" +: >"$OUT/ab.txt" +for i in $(seq 1 "$REPS"); do + for arm in A B; do + mode=""; [ "$arm" = "A" ] && mode="cpu" + log="$OUT/${arm}${i}.log" + ( + flock 9 || envfail "could not take the box lock $LOCK" + run_leg "${arm}${i}" "$mode" "$log" + ) 9>"$LOCK" + ms=$(decode_ms_per_token "$log") + echo "$arm $i $ms" >>"$OUT/ab.txt" + echo "${arm}${i} = $ms ms/token" + done +done + +awk ' + { v[$1""$2] = $3; arm[$1] += $3; n[$1]++ } + END { + a = arm["A"] / n["A"]; b = arm["B"] / n["B"]; + printf "A (CPU round trip) mean %.3f ms/token over %d\n", a, n["A"]; + printf "B (fused kernel) mean %.3f ms/token over %d\n", b, n["B"]; + printf "delta %.3f ms/token (%.2f%%)\n", a - b, 100.0 * (a - b) / a; + # Monotonic drift: same arm, first rep vs last rep. + if (n["A"] > 1) printf "drift within A %.2f%% (A1 %.3f -> A%d %.3f)\n", \ + 100.0 * (v["A" n["A"]] - v["A1"]) / v["A1"], v["A1"], n["A"], v["A" n["A"]]; + if (n["B"] > 1) printf "drift within B %.2f%% (B1 %.3f -> B%d %.3f)\n", \ + 100.0 * (v["B" n["B"]] - v["B1"]) / v["B1"], v["B1"], n["B"], v["B" n["B"]]; + print "REPORT THE DRIFT WITH THE DELTA. A delta smaller than the drift is not a result."; + }' "$OUT/ab.txt" + +echo "=== artifacts in $OUT ===" diff --git a/mistralrs-quant/kernels/arc_kvquant/arc_kvquant.cu b/mistralrs-quant/kernels/arc_kvquant/arc_kvquant.cu index 7afb89f6e..2f470e95e 100644 --- a/mistralrs-quant/kernels/arc_kvquant/arc_kvquant.cu +++ b/mistralrs-quant/kernels/arc_kvquant/arc_kvquant.cu @@ -36,11 +36,14 @@ // // 1. THIS CRATE COMPILES WITH `--use_fast_math` (mistralrs-quant/build.rs), so // a bare `a / b` here would be `div.approx.f32` while candle-kernels (no -// fast math) emits IEEE `div.rn.f32`, and `-ftz=true` would flush denormals -// candle keeps. Every float operation below is therefore an explicit -// `__f*_rn` intrinsic, which the CUDA Math API defines as IEEE-754 and -// unaffected by `-prec-div`/`-ftz`/`-fmad`. `fabsf`/`fmaxf` are avoided in -// favour of integer ops on the bit pattern for the same reason. +// fast math) emits IEEE `div.rn.f32`. The `__f*_rn` intrinsics are NOT a +// sufficient answer — checked against the emitted PTX, nvcc 13.1 rewrites +// them to `mul.rn.ftz.f32` / `div.rn.ftz.f32` / `add.rn.ftz.f32` under fast +// math, which flushes denormals candle keeps. Every float operation below is +// therefore INLINE PTX (`arc_fmul`/`arc_fadd`/`arc_fdiv`), which no +// optimisation flag rewrites; `fabsf`/`fmaxf` are avoided in favour of +// integer ops on the bit pattern for the same reason. The audit is one line: +// `nvcc -ptx ... | grep -c '\.ftz\.f32'` must be 0. // 2. The E4M3 rounding is a transcription of NVIDIA's // `__nv_cvt_double_to_fp8(x, __NV_SATFINITE, __NV_E4M3)` (cuda_fp8.hpp), // which is *also* what the Rust `float8` crate ports in `convert_to_fp8` @@ -72,6 +75,38 @@ namespace arc { // (inline `cvt.rn.bf16.f32`, immune to -ftz). // f32 -> f16 : likewise `__float2half` (inline `cvt.rn.f16.f32`). // --------------------------------------------------------------------------- +// --------------------------------------------------------------------------- +// IEEE f32 ops as inline PTX. +// +// MEASURED, not assumed: `--use_fast_math` (this crate's build.rs) rewrites the +// `__fmul_rn`/`__fadd_rn`/`__fdiv_rn` intrinsics to their `.ftz` forms — nvcc +// 13.1 emitted `mul.rn.ftz.f32` / `div.rn.ftz.f32` / `add.rn.ftz.f32` for them. +// The rounding mode survives, but denormal flush-to-zero does not match +// candle-kernels, which compiles WITHOUT fast math and emits `div.rn.f32`. +// +// The difference happens to be unobservable here (`scale` is floored at 1e-12, +// and E4M3's smallest denormal is 2^-9, so any f32 denormal is quantized to +// zero either way) — but bit-parity should not rest on a chain of reasoning +// about where denormals cannot reach. Inline PTX is never rewritten by +// optimisation flags, so these ops are IEEE non-ftz by construction, and +// `grep -c '\.ftz\.f32'` over the emitted PTX is a one-line audit of that. +// --------------------------------------------------------------------------- +__device__ __forceinline__ float arc_fmul(float a, float b) { + float r; + asm("mul.rn.f32 %0, %1, %2;" : "=f"(r) : "f"(a), "f"(b)); + return r; +} +__device__ __forceinline__ float arc_fadd(float a, float b) { + float r; + asm("add.rn.f32 %0, %1, %2;" : "=f"(r) : "f"(a), "f"(b)); + return r; +} +__device__ __forceinline__ float arc_fdiv(float a, float b) { + float r; + asm("div.rn.f32 %0, %1, %2;" : "=f"(r) : "f"(a), "f"(b)); + return r; +} + template __device__ __forceinline__ float arc_to_f32(T v); template <> __device__ __forceinline__ float arc_to_f32<__nv_bfloat16>(__nv_bfloat16 v) { return __bfloat162float(v); @@ -98,12 +133,12 @@ template <> __device__ __forceinline__ float arc_from_f32(float v) { retu // `affine(mul, add)` to `x * mul + add` with `mul`/`add` first narrowed to the // tensor dtype. `x * 1.0f` is exact, so the second affine is exactly an add, // and FMA contraction of `x * mul + 0.0f` is exactly the multiply. Both are -// therefore reproduced by one `__fmul_rn` and one `__fadd_rn`. +// therefore reproduced by one `arc_fmul` and one `arc_fadd` (inline PTX, see above). // --------------------------------------------------------------------------- __device__ __forceinline__ float arc_kv_block_scale(float amax) { const float inv_max = (float)(1.0 / ARC_KV_E4M3_MAX); const float eps = (float)1e-12; - return __fadd_rn(__fmul_rn(amax, inv_max), eps); + return arc_fadd(arc_fmul(amax, inv_max), eps); } // --------------------------------------------------------------------------- @@ -218,7 +253,7 @@ __global__ void arc_kv_fp8_quantize_kernel( for (int e = lane; e < block_w; e += ARC_KV_WARP) { const float v = arc_to_f32(krow[base + e]); - crow[base + e] = arc_f32_to_e4m3_code(__fdiv_rn(v, scale)); + crow[base + e] = arc_f32_to_e4m3_code(arc_fdiv(v, scale)); } // `amax` IS one of the block's own elements, so narrowing it back to the // activation dtype is exact - that is what lets dequant rebuild the @@ -270,7 +305,7 @@ __global__ void arc_kv_fp8_dequantize_kernel( const float amax = arc_to_f32(srow[rope_dim + blk]); const float scale = arc_kv_block_scale(amax); for (int e = lane; e < block_w; e += ARC_KV_WARP) { - orow[base + e] = arc_from_f32(__fmul_rn(lut[crow[base + e]], scale)); + orow[base + e] = arc_from_f32(arc_fmul(lut[crow[base + e]], scale)); } } @@ -419,7 +454,7 @@ __global__ void arc_kv_fp8_quantize_mutant_kernel( const float scale = arc_kv_block_scale(amax); for (int e = lane; e < block_w; e += ARC_KV_WARP) { const float v = arc_to_f32(krow[base + e]); - crow[base + e] = arc_f32_to_e4m3_code_truncating(__fdiv_rn(v, scale)); + crow[base + e] = arc_f32_to_e4m3_code_truncating(arc_fdiv(v, scale)); } if (lane == 0) { srow[rope_dim + blk] = arc_from_f32(amax); From 259369b2ee5dd9cad1dfb3c0249ac5596359f08c Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Tue, 18 Aug 2026 10:49:54 +0100 Subject: [PATCH 03/27] test(ArcKV/Fp8): exhaustive 2^32 E4M3 sweep + its negative control Runs on hardware with nvcc alone (no cargo build). Compares the kernel's transcription against NVIDIA's software reference (what the Rust float8 crate ports, hence what candle's CPU cast computes) and against the sm_90 hardware path, over ALL 2^32 f32 bit patterns. Measured on H200 / CUDA 13.1: visited 4,294,967,296 of 4,294,967,296 mismatch vs NVIDIA sw 0 mismatch sw vs hw 0 negative control 123,731,850 inputs caught (2.88%) The visited counter and the -DMUTANT=1 control exist because '0 mismatches' from a sweep that ran zero iterations is indistinguishable from a pass. Co-Authored-By: Claude Opus 5 (1M context) --- arc-tools/e4m3_exhaustive.cu | 137 +++++++++++++++++++++++++++++++++++ 1 file changed, 137 insertions(+) create mode 100644 arc-tools/e4m3_exhaustive.cu diff --git a/arc-tools/e4m3_exhaustive.cu b/arc-tools/e4m3_exhaustive.cu new file mode 100644 index 000000000..a93db1348 --- /dev/null +++ b/arc-tools/e4m3_exhaustive.cu @@ -0,0 +1,137 @@ +// EXHAUSTIVE proof that arc_kvquant.cu's E4M3 conversion IS NVIDIA's. +// +// Parent system: ArcQuant / TurboQuant (ArcInfer / ArcKV / Fp8). +// +// D33 — a check you have not seen fail is not a check. This sweeps ALL 2^32 +// f32 bit patterns (not a sample) and compares three implementations: +// +// arc::arc_f32_to_e4m3_code(x) the transcription the +// fused kernel ships +// __nv_cvt_double_to_fp8((double)x, SATFINITE, E4M3) NVIDIA's SOFTWARE +// reference — the routine +// the Rust `float8` crate +// ports as `convert_to_fp8`, +// hence exactly what +// `F8E4M3::from_f32`, and +// therefore candle's CPU +// cast, computes +// __nv_cvt_float_to_fp8(x, SATFINITE, E4M3) NVIDIA's HARDWARE path +// (cvt.rn.satfinite.e4m3x2.f32 +// on sm_89+) +// +// One mismatch against the software reference is a bug in the kernel. Software +// vs hardware is counted separately so the two can never be confused. +// +// Two guards make a green result mean something: +// * `visited` counts the values actually processed; a sweep that silently +// ran zero iterations would otherwise report "0 mismatches" and look +// identical to a pass. +// * `-DMUTANT=1` runs the same sweep against the deliberately-wrong +// truncating variant. If THAT also reports 0 mismatches the sweep is +// vacuous and the exit code says so. +// +// Build + run (needs only nvcc + a GPU, not the cargo build): +// nvcc -std=c++17 -O3 -arch=sm_90 --use_fast_math --expt-relaxed-constexpr \ +// -U__CUDA_NO_BFLOAT16_CONVERSIONS__ -DMUTANT=0 \ +// -I/mistralrs-quant/kernels/arc_kvquant e4m3_exhaustive.cu -o sweep0 +// Exit 0 pass, 1 mismatch, 2 environment/vacuous. +#ifndef MUTANT +#define MUTANT 0 +#endif + +#include +#include +#include +#include + +#include "arc_kvquant.cu" + +__global__ void sweep(unsigned long long base, unsigned long long n, + unsigned long long *mismatch_sw, + unsigned long long *mismatch_hw, unsigned *first_bad, + unsigned long long *visited) { + for (unsigned long long i = + blockIdx.x * (unsigned long long)blockDim.x + threadIdx.x; + i < n; i += (unsigned long long)gridDim.x * blockDim.x) { + const unsigned bits = (unsigned)(base + i); + const float x = __uint_as_float(bits); + const uint8_t mine = MUTANT ? arc::arc_f32_to_e4m3_code_truncating(x) + : arc::arc_f32_to_e4m3_code(x); + const uint8_t sw = + (uint8_t)__nv_cvt_double_to_fp8((double)x, __NV_SATFINITE, __NV_E4M3); + const uint8_t hw = + (uint8_t)__nv_cvt_float_to_fp8(x, __NV_SATFINITE, __NV_E4M3); + if (mine != sw) { + if (atomicAdd(mismatch_sw, 1ULL) == 0ULL) { + *first_bad = bits; + } + } + if (sw != hw) { + atomicAdd(mismatch_hw, 1ULL); + } + atomicAdd(visited, 1ULL); + } +} + +int main() { + unsigned long long *d_sw, *d_hw, *d_vis; + unsigned *d_first; + if (cudaMalloc(&d_sw, 8) != cudaSuccess || cudaMalloc(&d_hw, 8) != cudaSuccess || + cudaMalloc(&d_vis, 8) != cudaSuccess || + cudaMalloc(&d_first, 4) != cudaSuccess) { + printf("FATAL cudaMalloc\n"); + return 2; + } + cudaMemset(d_sw, 0, 8); + cudaMemset(d_hw, 0, 8); + cudaMemset(d_vis, 0, 8); + cudaMemset(d_first, 0, 4); + + const unsigned long long TOTAL = 1ULL << 32; + const unsigned long long CHUNK = 1ULL << 28; + for (unsigned long long base = 0; base < TOTAL; base += CHUNK) { + sweep<<<8192, 256>>>(base, CHUNK, d_sw, d_hw, d_first, d_vis); + cudaError_t e = cudaDeviceSynchronize(); + if (e != cudaSuccess) { + printf("FATAL cuda: %s\n", cudaGetErrorString(e)); + return 2; + } + } + + unsigned long long sw = 0, hw = 0, vis = 0; + unsigned first = 0; + cudaMemcpy(&sw, d_sw, 8, cudaMemcpyDeviceToHost); + cudaMemcpy(&hw, d_hw, 8, cudaMemcpyDeviceToHost); + cudaMemcpy(&vis, d_vis, 8, cudaMemcpyDeviceToHost); + cudaMemcpy(&first, d_first, 4, cudaMemcpyDeviceToHost); + + printf("MUTANT %d\n", MUTANT); + printf("visited %llu of %llu f32 bit patterns\n", vis, TOTAL); + if (vis != TOTAL) { + printf("FATAL the sweep did not visit every value; a 0 here means nothing\n"); + return 2; + } + printf("mismatch vs NVIDIA sw %llu\n", sw); + printf("mismatch sw vs hw %llu\n", hw); + + if (MUTANT) { + if (sw == 0) { + printf("FATAL VACUOUS: the deliberately-wrong kernel also matched on " + "every input, so this sweep cannot fail and proves nothing\n"); + return 2; + } + printf("NEGATIVE CONTROL OK: the wrong kernel is caught on %llu inputs " + "(%.2f%%)\n", + sw, 100.0 * (double)sw / (double)TOTAL); + return 0; + } + if (sw) { + float x; + memcpy(&x, &first, 4); + printf("first bad input 0x%08x (%.9g)\n", first, x); + return 1; + } + printf("RESULT: arc_f32_to_e4m3_code is bit-identical to NVIDIA's E4M3 " + "conversion on every one of the 2^32 f32 values.\n"); + return 0; +} From 112cbb4e1a37fbbe00fa1f2fa43b90c9eda18beb Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Tue, 18 Aug 2026 10:59:40 +0100 Subject: [PATCH 04/27] test(ArcLab): per-step D2H + launch counter, validated on a known answer The 43 blocking cuMemcpyDtoHAsync_v2 per V4 decode step are invisible to the obvious instrument (*Synchronize* is 0.0 calls/step), so this counts the copies themselves and pins the step count from the trace with two independent anchors rather than assuming it. Known-answer test on the recorded baseline trace (/root/budget-chain/nsys): D2H PER STEP 44.09 (BUDGET_V4_B1.md records 44) LAUNCHES PER STEP 9131.5 (records 9,131) 1,792 B x 8,144 = 43.09/step; 517,120 B x 189 = 1.00/step -- the exact two DtoH sizes the budget names. Instrument validated before being pointed at the new traces. Drops the earlier measure_kv_fp8_fused.sh: it was written before the box paths were known and its exclusivity check had a defect (it cleared the pid it was meant to compare against). Replaced by the scripts actually run on the box. Co-Authored-By: Claude Opus 5 (1M context) --- arc-tools/kv_fp8_count_per_step.py | 102 +++++++++++++ arc-tools/measure_kv_fp8_fused.sh | 223 ----------------------------- 2 files changed, 102 insertions(+), 223 deletions(-) create mode 100644 arc-tools/kv_fp8_count_per_step.py delete mode 100755 arc-tools/measure_kv_fp8_fused.sh diff --git a/arc-tools/kv_fp8_count_per_step.py b/arc-tools/kv_fp8_count_per_step.py new file mode 100644 index 000000000..224945b22 --- /dev/null +++ b/arc-tools/kv_fp8_count_per_step.py @@ -0,0 +1,102 @@ +#!/usr/bin/env python3 +"""Per-decode-step D2H and kernel-launch counts from an nsys sqlite. + +Parent system: ArcLab (measuring ArcInfer / ArcKV / Fp8). + +WHY THIS EXISTS + The 43 blocking `cuMemcpyDtoHAsync_v2` per V4 decode step are invisible to + the obvious instrument: `*Synchronize*` is 0.0 calls/step in this workload, + so anyone counting `cudaStreamSynchronize` concludes "no syncs" and is + wrong. This counts the copies themselves. + +THE STEP COUNT IS PINNED FROM THE TRACE, NOT ASSUMED + A per-step number divided by a guessed step count is a fabricated number. + Two independent anchors must agree: + A. the once-per-step logits readback -- the largest DtoH size, whose + occurrence count IS the number of steps in the window; + B. structural agreement -- kernels whose counts are integer multiples of + that same step count (the 43-layer loop makes many). Fewer than three + such kernels means the window is not steady-state decode and the + script exits 2. + +Exits 2 on environment/analysis failure, never 1. + +Usage: kv_fp8_count_per_step.py [tail_seconds] +""" +import collections +import sqlite3 +import sys + + +def die(msg): + print(f"FATAL {msg}", file=sys.stderr) + sys.exit(2) + + +def main(): + if len(sys.argv) < 2: + die("usage: kv_fp8_count_per_step.py [tail_seconds]") + tail_s = float(sys.argv[2]) if len(sys.argv) > 2 else 20.0 + + db = sqlite3.connect(sys.argv[1]) + cur = db.cursor() + strings = {i: v for i, v in cur.execute("SELECT id, value FROM StringIds")} + KT = "CUPTI_ACTIVITY_KIND_KERNEL" + MT = "CUPTI_ACTIVITY_KIND_MEMCPY" + + end = cur.execute(f"SELECT MAX(end) FROM {KT}").fetchone()[0] + if end is None: + die("no kernels in the trace") + win0 = end - int(tail_s * 1e9) + + # copyKind 2 == device-to-host. + sizes = list(cur.execute( + f"SELECT bytes, COUNT(*) FROM {MT} WHERE start>={win0} AND copyKind=2 " + f"GROUP BY bytes ORDER BY COUNT(*) DESC")) + if not sizes: + die("no DtoH copies in the window") + total_d2h = sum(c for _, c in sizes) + + # ---- Anchor A: the logits readback, once per decode step. + big = [(b, c) for b, c in sizes if b > 100_000] + if not big: + die("no large (>100 kB) DtoH found; cannot pin the step count") + steps = max(c for _, c in big) + if steps < 20: + die(f"only {steps} steps in the window; widen tail_seconds") + + kc = collections.Counter() + for (n,) in cur.execute(f"SELECT shortName FROM {KT} WHERE start>={win0}"): + kc[strings.get(n, str(n))] += 1 + if not kc: + die("no kernels in the window") + + # ---- Anchor B: kernels whose counts are integer multiples of `steps`. + agreeing = [] + for name, c in kc.items(): + k = c / steps + if k >= 0.98 and abs(k - round(k)) <= 0.02 * max(round(k), 1): + agreeing.append((name, c, round(k))) + if len(agreeing) < 3: + die(f"step count {steps} unconfirmed: only {len(agreeing)} kernels are " + f"integer multiples of it; the window is probably not steady-state " + f"decode") + + total_kern = sum(kc.values()) + print(f"window last {tail_s:.0f} s") + print(f"steps (logits anchor) {steps}") + print(f"steps confirmed by {len(agreeing)} kernels at integer multiples") + print(f"D2H total {total_d2h}") + print(f"D2H PER STEP {total_d2h / steps:.2f}") + print(f"kernels total {total_kern}") + print(f"LAUNCHES PER STEP {total_kern / steps:.1f}") + print("DtoH by size:") + for b, c in sizes[:8]: + print(f" {b:>12,d} B x {c:>9,d} {c / steps:8.2f}/step") + print("anchor-B kernels (count = k x steps):") + for name, c, k in sorted(agreeing, key=lambda t: -t[1])[:6]: + print(f" {c:>9,d} = {k:>4d} x steps {name[:52]}") + + +if __name__ == "__main__": + main() diff --git a/arc-tools/measure_kv_fp8_fused.sh b/arc-tools/measure_kv_fp8_fused.sh deleted file mode 100755 index fcf03815a..000000000 --- a/arc-tools/measure_kv_fp8_fused.sh +++ /dev/null @@ -1,223 +0,0 @@ -#!/usr/bin/env bash -# ArcKV/Fp8 — measure the fused E4M3 quantize+dequantize kernel against the -# CPU round trip it replaces. Runs ON THE BOX. Nothing here is a proxy: every -# number is produced by the real binary on the real model. -# -# WHAT IT PRODUCES (the four numbers the change is judged on) -# 1. `cuMemcpyDtoHAsync_v2` calls per decode step, before and after. -# The 44/step are the thing being removed and the reason CUDA graph -# capture is impossible. NOTE `*Synchronize*` is 0.0 calls/step in this -# workload — counting `cudaStreamSynchronize` reports "no syncs" and is -# wrong. Count the D2H. -# 2. Kernel launches per decode step, before and after. Op count is the -# disease (9,131/token, median kernel 1.18 us), so this matters more than -# wall time. -# 3. Interleaved A-B-A-B ms/token with the monotonic drift stated. This box -# showed 3.3% drift across one run — larger than a real arm difference — -# so A-then-B is a fabricated comparison. -# 4. Bit-parity: the GPU-gated `kv_fp8_fused_is_bit_identical_to_cpu_exact` -# test, plus its negative control. -# -# RULES ENFORCED HERE, NOT ASSUMED -# * The box is shared with the ArcGraph chain. Every timing leg holds -# /root/.arc-bench.lock under flock; a contended number is a fabricated -# number. -# * Exclusivity is asserted BEFORE AND AFTER every leg from -# `nvidia-smi --query-compute-apps` (the lock file is not an occupancy -# signal — an 87 GB server has been resident with no lock at all, and a -# 77 GB server has run with the lock reading FREE). A neighbour appearing -# mid-leg aborts the run. -# * Environment failure exits 2. A failed measurement exits 1. They are not -# the same thing and must never be reported as the same thing. -# * Engagement is asserted before any null/neutral result is believed: if -# the arm that is supposed to use the fused kernel shows the same D2H -# count as the arm that is not, the run aborts rather than reporting "no -# difference". -set -uo pipefail - -REPO="${REPO:-/root/arc-wt/fp8}" -MODEL="${MODEL:-deepseek-ai/DeepSeek-V4-Flash}" -ARCH="${ARCH:-deepseekv4}" -ISQ="${ISQ:-qtip2}" -PROMPT_LEN="${PROMPT_LEN:-128}" -GEN_LEN="${GEN_LEN:-64}" -MAX_SEQ_LEN="${MAX_SEQ_LEN:-1024}" -OUT="${OUT:-/root/kvfp8-fused}" -LOCK="${LOCK:-/root/locks/bench.lock}" -BIN="${BIN:-${CARGO_TARGET_DIR:-$REPO/target}/release/mistralrs}" -REPS="${REPS:-2}" # A-B-A-B => REPS=2 - -mkdir -p "$OUT" - -envfail() { - echo "ENVFAIL: $*" >&2 - exit 2 -} -fail() { - echo "FAIL: $*" >&2 - exit 1 -} - -command -v nvidia-smi >/dev/null 2>&1 || envfail "no nvidia-smi" -command -v nsys >/dev/null 2>&1 || echo "WARN: no nsys; the D2H/launch counts will be skipped" >&2 -[ -x "$BIN" ] || envfail "no mistralrs binary at $BIN (build with --features 'cuda flash-attn')" - -# --------------------------------------------------------------------------- -# Exclusivity. Bracket every leg: a V4 load shows near-zero VRAM for most of a -# minute, so "compute-apps empty" sampled during a neighbour's load is -# indistinguishable from an idle box. -# --------------------------------------------------------------------------- -MYPID="" -assert_exclusive() { - local where="$1" - local apps - apps=$(nvidia-smi --query-compute-apps=pid --format=csv,noheader | tr -d ' ' | sort -u | tr '\n' ',') - apps="${apps%,}" - if [ -n "$MYPID" ]; then - [ "$apps" = "$MYPID" ] || envfail "[$where] compute-apps=[$apps], expected only $MYPID" - else - [ -z "$apps" ] || envfail "[$where] box not idle: compute-apps=[$apps]" - fi -} - -# --------------------------------------------------------------------------- -# One measured leg. $1 = arm label, $2 = ARC_KV_FP8_MODE value ("" => default, -# which is the fused kernel), $3 = log path, $4 = "nsys" to trace. -# --------------------------------------------------------------------------- -run_leg() { - local arm="$1" mode="$2" log="$3" trace="${4:-}" - local -a cmd=("$BIN" bench -m "$MODEL" -a "$ARCH" --isq "$ISQ" - --prompt-len "$PROMPT_LEN" --gen-len "$GEN_LEN" --max-seq-len "$MAX_SEQ_LEN") - - assert_exclusive "pre-$arm" - ( - if [ -n "$mode" ]; then export ARC_KV_FP8_MODE="$mode"; else unset ARC_KV_FP8_MODE; fi - unset ARC_GPU_ACT_QUANT - if [ "$trace" = "nsys" ]; then - nsys profile -t cuda -o "$OUT/$arm" --force-overwrite true \ - --cuda-memory-usage false "${cmd[@]}" - else - "${cmd[@]}" - fi - ) >"$log" 2>&1 & - local pid=$! - # `pgrep -f`/`pkill -f` would match this script's own command line; use the - # job pid we already hold instead. - MYPID="" - wait "$pid" - local rc=$? - MYPID="" - assert_exclusive "post-$arm" - [ $rc -eq 0 ] || fail "$arm leg exited $rc; see $log" -} - -# --------------------------------------------------------------------------- -# ms/token out of a bench log. Fails loudly rather than emitting an empty -# string that would silently become a "0.0 ms/token improvement". -# --------------------------------------------------------------------------- -decode_ms_per_token() { - local log="$1" - local tps - tps=$(grep -oE 'tok_per_s_decode[^0-9]*[0-9]+\.[0-9]+' "$log" | tail -1 | - grep -oE '[0-9]+\.[0-9]+$') - [ -n "$tps" ] || tps=$(grep -oiE 'decode[^0-9]*([0-9]+\.[0-9]+) *tok' "$log" | tail -1 | - grep -oE '[0-9]+\.[0-9]+') - [ -n "$tps" ] || fail "no decode throughput in $log — do not report a delta from a missing number" - awk -v t="$tps" 'BEGIN { printf "%.3f", 1000.0 / t }' -} - -# --------------------------------------------------------------------------- -# 0. Bit-parity + its negative control. Cheap, and it gates everything else: -# a faster kernel that stores different bytes is not a faster kernel. -# --------------------------------------------------------------------------- -echo "=== 0. bit-parity (GPU-gated test + D33 negative control) ===" -( - cd "$REPO" || envfail "no repo at $REPO" - cargo test -p mistralrs-core --features "cuda flash-attn" --lib \ - kv_fp8_fused_is_bit_identical_to_cpu_exact -- --nocapture --exact \ - models::dsv4_kv_fp8::tests::kv_fp8_fused_is_bit_identical_to_cpu_exact -) 2>&1 | tee "$OUT/parity.log" -prc=${PIPESTATUS[0]} -# `set -e` does not survive a pipe; the status is read explicitly. -[ "$prc" -eq 2 ] && envfail "parity test could not find a CUDA device" -[ "$prc" -eq 0 ] || fail "bit-parity FAILED — the fused kernel does not store what the CPU path stores" -grep -q "test result: ok. 1 passed" "$OUT/parity.log" || - fail "parity test reported no results — 'no failures' is not 'ran'" -echo "bit-parity: PASS" - -# --------------------------------------------------------------------------- -# 1 + 2. D2H count and launch count per decode step, both arms, under nsys. -# --------------------------------------------------------------------------- -if command -v nsys >/dev/null 2>&1; then - echo "=== 1+2. nsys: DtoH copies and kernel launches per step ===" - flock "$LOCK" bash -c "$(declare -f run_leg assert_exclusive envfail fail); \ - OUT='$OUT' BIN='$BIN' MODEL='$MODEL' ARCH='$ARCH' ISQ='$ISQ' \ - PROMPT_LEN='$PROMPT_LEN' GEN_LEN='$GEN_LEN' MAX_SEQ_LEN='$MAX_SEQ_LEN' \ - bash -c 'true'" || true - - for arm in before after; do - mode=""; [ "$arm" = "before" ] && mode="cpu" - ( - flock 9 || envfail "could not take the box lock $LOCK" - run_leg "$arm" "$mode" "$OUT/$arm.nsys.log" nsys - ) 9>"$LOCK" - nsys stats --report cuda_api_sum --format csv "$OUT/$arm.nsys-rep" \ - >"$OUT/$arm.api.csv" 2>"$OUT/$arm.api.err" || - echo "WARN: nsys stats cuda_api_sum failed for $arm" >&2 - nsys stats --report cuda_gpu_kern_sum --format csv "$OUT/$arm.nsys-rep" \ - >"$OUT/$arm.kern.csv" 2>>"$OUT/$arm.api.err" || - echo "WARN: nsys stats cuda_gpu_kern_sum failed for $arm" >&2 - - d2h=$(awk -F, '/cuMemcpyDtoHAsync_v2/ { gsub(/"/,"",$3); s+=$3 } END { print s+0 }' "$OUT/$arm.api.csv") - launches=$(awk -F, 'NR>1 { gsub(/"/,"",$3); s+=$3 } END { print s+0 }' "$OUT/$arm.kern.csv") - echo "$arm: cuMemcpyDtoHAsync_v2=$d2h kernel_launches=$launches (steps=$GEN_LEN)" - echo "$arm $d2h $launches" >>"$OUT/counts.txt" - done - - # Engagement (D18): the two arms MUST differ in D2H count. Equal counts mean - # the fused arm never engaged, and every later number is the same code twice. - b=$(awk '$1=="before"{print $2}' "$OUT/counts.txt") - a=$(awk '$1=="after"{print $2}' "$OUT/counts.txt") - [ -n "$b" ] && [ -n "$a" ] || fail "missing D2H counts — no results is not a null result" - [ "$a" -lt "$b" ] || fail "ENGAGEMENT: after-arm D2H ($a) is not below before-arm ($b); the fused kernel did not run" - awk -v b="$b" -v a="$a" -v n="$GEN_LEN" 'BEGIN { - printf "D2H per step: before %.2f -> after %.2f (removed %.2f/step)\n", b/n, a/n, (b-a)/n }' -fi - -# --------------------------------------------------------------------------- -# 3. Interleaved A-B-A-B ms/token. A = CPU round trip, B = fused kernel. -# The drift is reported alongside the delta, because on this box a 3.3% -# monotonic drift once exceeded the arm difference. -# --------------------------------------------------------------------------- -echo "=== 3. interleaved A-B-A-B ms/token ===" -: >"$OUT/ab.txt" -for i in $(seq 1 "$REPS"); do - for arm in A B; do - mode=""; [ "$arm" = "A" ] && mode="cpu" - log="$OUT/${arm}${i}.log" - ( - flock 9 || envfail "could not take the box lock $LOCK" - run_leg "${arm}${i}" "$mode" "$log" - ) 9>"$LOCK" - ms=$(decode_ms_per_token "$log") - echo "$arm $i $ms" >>"$OUT/ab.txt" - echo "${arm}${i} = $ms ms/token" - done -done - -awk ' - { v[$1""$2] = $3; arm[$1] += $3; n[$1]++ } - END { - a = arm["A"] / n["A"]; b = arm["B"] / n["B"]; - printf "A (CPU round trip) mean %.3f ms/token over %d\n", a, n["A"]; - printf "B (fused kernel) mean %.3f ms/token over %d\n", b, n["B"]; - printf "delta %.3f ms/token (%.2f%%)\n", a - b, 100.0 * (a - b) / a; - # Monotonic drift: same arm, first rep vs last rep. - if (n["A"] > 1) printf "drift within A %.2f%% (A1 %.3f -> A%d %.3f)\n", \ - 100.0 * (v["A" n["A"]] - v["A1"]) / v["A1"], v["A1"], n["A"], v["A" n["A"]]; - if (n["B"] > 1) printf "drift within B %.2f%% (B1 %.3f -> B%d %.3f)\n", \ - 100.0 * (v["B" n["B"]] - v["B1"]) / v["B1"], v["B1"], n["B"], v["B" n["B"]]; - print "REPORT THE DRIFT WITH THE DELTA. A delta smaller than the drift is not a result."; - }' "$OUT/ab.txt" - -echo "=== artifacts in $OUT ===" From 41a9343dfd891b0e0c6580a468f87e107d1a8006 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Tue, 18 Aug 2026 11:03:25 +0100 Subject: [PATCH 05/27] test(ArcKV/Fp8): nsys count harness with lock + VRAM gate MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Counts survive contention (nsys traces only this process), but V4 is ~79 GB of the H200's 143 GB, so holding the bench lock is not enough — the previous holder's process can still be resident. Gates on nvidia-smi free memory and exits 2 on OOM or a missing report rather than reporting a partial trace. Co-Authored-By: Claude Opus 5 (1M context) --- arc-tools/kv_fp8_nsys_ab.sh | 67 +++++++++++++++++++++++++++++++++++++ 1 file changed, 67 insertions(+) create mode 100755 arc-tools/kv_fp8_nsys_ab.sh diff --git a/arc-tools/kv_fp8_nsys_ab.sh b/arc-tools/kv_fp8_nsys_ab.sh new file mode 100755 index 000000000..6d4d40fbe --- /dev/null +++ b/arc-tools/kv_fp8_nsys_ab.sh @@ -0,0 +1,67 @@ +#!/bin/bash +# ArcKV/Fp8 — nsys legs: `cuMemcpyDtoHAsync_v2` copies and kernel launches PER +# DECODE STEP, before (CPU round trip) and after (fused device kernel). +# +# These are COUNTS. Box contention changes wall-clock, not how many D2H calls a +# decode step makes, so a contended run still yields a valid count — but the +# card must physically have room: V4 is ~79 GB of the H200's 143 GB, so two of +# these do not fit, and an OOM mid-trace is an environment failure (exit 2), not +# a result. +# +# Counted with arc-tools/kv_fp8_count_per_step.py, which pins the step count +# from the trace and was first validated against the recorded baseline (it +# reproduces 44.09 D2H/step and 9,131.5 launches/step). +set -u +OUT=${OUT:-/root/kvfp8} +BASE=${BASE:-/root/budget-chain} +BIN=${BIN:-/root/arc-wt/fp8-target/release/mistralrs} +DUR=${DUR:-280} +GEN=${GEN:-3000} +export PATH=/usr/local/cuda-13.1/bin:$PATH +export LD_LIBRARY_PATH=/usr/local/cuda/compat:${LD_LIBRARY_PATH:-} +unset ARC_TIME_DECODE V4_STATS V4_NAN_DEBUG V4_TRACE ARC_PROFILE ARC_GPU_ACT_QUANT +[ -x "$BIN" ] || { echo "FATAL_NO_BINARY $BIN"; exit 2; } +command -v nsys >/dev/null || { echo "FATAL_NO_NSYS"; exit 2; } + +# Holding the bench lock is not the same as the previous holder's process having +# exited; gate on what the card actually reports. +wait_for_vram() { + local need=${1:-100000} i=0 free + while :; do + free=$(nvidia-smi --query-gpu=memory.free --format=csv,noheader,nounits) + [ "${free:-0}" -ge "$need" ] && { echo "vram ok: ${free} MiB free"; return 0; } + i=$((i + 1)) + [ "$i" -gt 240 ] && { echo "FATAL_NO_VRAM free=${free}MiB after 20 min"; exit 2; } + sleep 5 + done +} + +for arm in before:cpu after:fused; do + A=${arm%%:*} + M=${arm##*:} + echo "=== leg $A (ARC_KV_FP8_MODE=$M) $(date -u +%T) ===" + wait_for_vram 100000 + rm -f "$OUT/$A".nsys-rep "$OUT/$A".sqlite + ARC_KV_FP8_MODE=$M nsys profile --trace=cuda --sample=none --cpuctxsw=none \ + --cuda-memory-usage=false --duration="$DUR" --kill=sigterm \ + --force-overwrite=true --output="$OUT/$A" \ + "$BIN" bench -m "$BASE/src" -a deepseekv4 \ + --from-uqff "$BASE/uqff/qtip2-0.uqff" --max-seqs 1 --prefix-cache-n 0 \ + --prompt-len 64 --gen-len "$GEN" --iterations 1 --warmup 0 \ + >"$OUT/$A.nsys.log" 2>&1 + RC=$? + echo "NSYS_RC=$RC" + # 143 = SIGTERM, which is how --kill ends the window and is EXPECTED. + [ "$RC" = "139" ] && { echo "FATAL_TARGET_SIGSEGV"; exit 2; } + grep -qi "out of memory" "$OUT/$A.nsys.log" && { echo "FATAL_OOM $A"; exit 2; } + [ -f "$OUT/$A.nsys-rep" ] || { echo "FATAL_NO_REPORT $A"; exit 2; } + nsys stats --force-export=true --report cuda_api_sum --format csv \ + --output "$OUT/${A}_api" "$OUT/$A.nsys-rep" >/dev/null 2>&1 + [ -f "$OUT/$A.sqlite" ] || { echo "FATAL_NO_SQLITE $A"; exit 2; } +done + +echo "=== ANALYSIS ===" +for A in before after; do + echo "--- $A ---" + python3 "$OUT/count.py" "$OUT/$A.sqlite" 20 || exit 2 +done From ad75bdc29ace43607d707a744a4474845b0c1fc7 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Tue, 18 Aug 2026 11:06:20 +0100 Subject: [PATCH 06/27] test(ArcLab): count CUDA API calls per step too, incl. the invisible sync MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds cuMemcpyDtoHAsync_v2 / cuLaunchKernel / cuMemAllocAsync / cuMemFreeAsync per step, and prints ALL *Synchronize* alongside — it reads 0.00/step, which is exactly why counting syncs on this workload finds nothing and concludes wrongly. Known-answer re-test on the recorded baseline trace now reproduces every headline in BUDGET_V4_B1.md: 44.09/step cuMemcpyDtoHAsync_v2 (recorded 44) 11436.33 + 11436.23 alloc/free/step (recorded 11,436 each) 2818.31/step cuMemcpyHtoDAsync_v2 (recorded 2,818) 0.00/step ALL *Synchronize* (recorded 0.0) 9131.5 launches/step (recorded 9,131) Co-Authored-By: Claude Opus 5 (1M context) --- arc-tools/kv_fp8_count_per_step.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/arc-tools/kv_fp8_count_per_step.py b/arc-tools/kv_fp8_count_per_step.py index 224945b22..39bb43314 100644 --- a/arc-tools/kv_fp8_count_per_step.py +++ b/arc-tools/kv_fp8_count_per_step.py @@ -97,6 +97,29 @@ def main(): for name, c, k in sorted(agreeing, key=lambda t: -t[1])[:6]: print(f" {c:>9,d} = {k:>4d} x steps {name[:52]}") + # ---- CUDA API calls per step. `cuMemcpyDtoHAsync_v2` is the one that + # matters here and it is NOT visible as a "sync": `*Synchronize*` is 0.0 + # calls/step in this workload. An "async" copy costing ~109 us of host time + # is a pageable-memory staged copy, i.e. blocking, and a CUDA graph cannot + # record one. + api = collections.Counter() + try: + for (n,) in cur.execute( + "SELECT nameId FROM CUPTI_ACTIVITY_KIND_RUNTIME WHERE start>=?", + (win0,)): + api[strings.get(n, str(n))] += 1 + except sqlite3.Error: + print("(no CUPTI_ACTIVITY_KIND_RUNTIME table; API counts unavailable)") + return + watch = ("cuMemcpyDtoHAsync_v2", "cuLaunchKernel", "cuMemAllocAsync", + "cuMemFreeAsync", "cuMemcpyHtoDAsync_v2", "cuMemsetD8Async") + print("CUDA API per step:") + for name in watch: + print(f" {api.get(name, 0) / steps:10.2f}/step {name}") + syncs = sum(c for n, c in api.items() if "Synchronize" in n) + print(f" {syncs / steps:10.2f}/step ALL *Synchronize* " + f"(this is why counting syncs finds nothing)") + if __name__ == "__main__": main() From 02c38946b4cad0c750ddaa7d0f9fe24cf67803cb Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Tue, 18 Aug 2026 11:38:48 +0100 Subject: [PATCH 07/27] fix(ArcLab): hold the bench lock across GPU work ONLY, not report export MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Measured on the box: a chain wrapping its whole pipeline in one flock held /root/locks/bench.lock with 13-17 waiters queued while the GPU read 0 %, 0 MiB, 78.28 W. nsys report export and the per-step counting are pure CPU on files already written, and must not hold the card. The script now takes the lock itself, per leg, around the VRAM wait and the traced run, and releases it the instant the bench exits — and says so in the output so the release time is auditable. Do not wrap it in flock. Co-Authored-By: Claude Opus 5 (1M context) --- arc-tools/kv_fp8_nsys_ab.sh | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/arc-tools/kv_fp8_nsys_ab.sh b/arc-tools/kv_fp8_nsys_ab.sh index 6d4d40fbe..362a2b946 100755 --- a/arc-tools/kv_fp8_nsys_ab.sh +++ b/arc-tools/kv_fp8_nsys_ab.sh @@ -11,11 +11,20 @@ # Counted with arc-tools/kv_fp8_count_per_step.py, which pins the step count # from the trace and was first validated against the recorded baseline (it # reproduces 44.09 D2H/step and 9,131.5 launches/step). +# +# LOCK DISCIPLINE — this script takes /root/locks/bench.lock ITSELF, per leg, +# and holds it ONLY across the traced run. Do NOT wrap this script in `flock`: +# that holds the card through nsys report export and the counting, both of which +# are pure CPU on files already written. Measured on this box, a chain doing +# exactly that held the lock with 13-17 waiters queued while the GPU read +# `0 %, 0 MiB, 78.28 W`. set -u OUT=${OUT:-/root/kvfp8} BASE=${BASE:-/root/budget-chain} BIN=${BIN:-/root/arc-wt/fp8-target/release/mistralrs} DUR=${DUR:-280} +LOCK=${LOCK:-/root/locks/bench.lock} +LOCKWAIT=${LOCKWAIT:-5400} GEN=${GEN:-3000} export PATH=/usr/local/cuda-13.1/bin:$PATH export LD_LIBRARY_PATH=/usr/local/cuda/compat:${LD_LIBRARY_PATH:-} @@ -36,12 +45,16 @@ wait_for_vram() { done } +exec 9>"$LOCK" for arm in before:cpu after:fused; do A=${arm%%:*} M=${arm##*:} echo "=== leg $A (ARC_KV_FP8_MODE=$M) $(date -u +%T) ===" - wait_for_vram 100000 rm -f "$OUT/$A".nsys-rep "$OUT/$A".sqlite + # LOCK HELD ONLY HERE: VRAM wait + the traced run. Released the instant the + # bench exits, before any report export. + flock -w "$LOCKWAIT" 9 || { echo "FATAL_LOCK_TIMEOUT after ${LOCKWAIT}s"; exit 2; } + wait_for_vram 100000 ARC_KV_FP8_MODE=$M nsys profile --trace=cuda --sample=none --cpuctxsw=none \ --cuda-memory-usage=false --duration="$DUR" --kill=sigterm \ --force-overwrite=true --output="$OUT/$A" \ @@ -50,7 +63,8 @@ for arm in before:cpu after:fused; do --prompt-len 64 --gen-len "$GEN" --iterations 1 --warmup 0 \ >"$OUT/$A.nsys.log" 2>&1 RC=$? - echo "NSYS_RC=$RC" + flock -u 9 + echo "NSYS_RC=$RC (lock released $(date -u +%T); export below is CPU-only)" # 143 = SIGTERM, which is how --kill ends the window and is EXPECTED. [ "$RC" = "139" ] && { echo "FATAL_TARGET_SIGSEGV"; exit 2; } grep -qi "out of memory" "$OUT/$A.nsys.log" && { echo "FATAL_OOM $A"; exit 2; } From e57b95933ac60441ab0eb5e89a4f61791ce1d56c Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Tue, 18 Aug 2026 11:50:00 +0100 Subject: [PATCH 08/27] test(ArcKV/Fp8): interleaved A/B driver with per-leg lock + Xid check MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Lock covers exactly one leg (model load, which allocates, plus the timed run) and is released before parsing. Interleaves A-B-A-B and prints the within-arm drift beside the delta, because this box once drifted 3.3% monotonically — more than the arm difference. On a bench failure it dumps 'dmesg | grep -i xid' first: this box carries ~1,485 Xid GPU faults (ECC uncorrected 0, so not memory corruption), and a process killed by one dies with no error line. Distinguishing the box's fault from the code's is the difference between fixing a bug and chasing one that does not exist. Co-Authored-By: Claude Opus 5 (1M context) --- arc-tools/kv_fp8_ab.sh | 83 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 83 insertions(+) create mode 100755 arc-tools/kv_fp8_ab.sh diff --git a/arc-tools/kv_fp8_ab.sh b/arc-tools/kv_fp8_ab.sh new file mode 100755 index 000000000..876d5042c --- /dev/null +++ b/arc-tools/kv_fp8_ab.sh @@ -0,0 +1,83 @@ +#!/usr/bin/env bash +# ArcKV/Fp8 — interleaved A-B-A-B decode timing for the fused E4M3 kernel. +# A = ARC_KV_FP8_MODE=cpu (43 blocking D2H per step) +# B = ARC_KV_FP8_MODE=fused (one device kernel, 0 D2H) +# +# A-then-B is banned on this box: a 3.3% monotonic drift once exceeded the arm +# difference. The drift within each arm is printed next to the delta, and a +# delta smaller than the drift is not a result. +# +# LOCK DISCIPLINE. The lock covers exactly one thing: a leg (model load, which +# allocates GPU memory, plus the timed run). It is taken immediately before and +# released immediately after — NOT held across parsing, and NOT wrapped around +# the whole script. Six chains share this H200; one holding the lock through +# post-processing starved the trellis chain for 40 minutes while the card read +# 0 %, 0 MiB, 78 W. +set -u +OUT=${OUT:-/root/kvfp8} +BASE=${BASE:-/root/budget-chain} +BIN=${BIN:-/root/arc-wt/fp8-target/release/mistralrs} +LOCK=${LOCK:-/root/locks/bench.lock} +LOCKWAIT=${LOCKWAIT:-3600} +REPS=${REPS:-2} +GEN=${GEN:-260} +export PATH=/usr/local/cuda-13.1/bin:$PATH +export LD_LIBRARY_PATH=/usr/local/cuda/compat:${LD_LIBRARY_PATH:-} +unset ARC_TIME_DECODE V4_STATS V4_NAN_DEBUG V4_TRACE ARC_PROFILE ARC_GPU_ACT_QUANT +[ -x "$BIN" ] || { echo "FATAL_NO_BINARY $BIN"; exit 2; } + +exec 9>"$LOCK" + +leg() { # $1 = tag, $2 = mode + local tag=$1 mode=$2 log="$OUT/$1.log" free rc ms + flock -w "$LOCKWAIT" 9 || { echo "FATAL_LOCK_TIMEOUT ${LOCKWAIT}s"; exit 2; } + # Holding the lock is not the same as the card being free: V4 is ~79 GB of + # 143 GB and the previous holder's process can outlive its script. + for _ in $(seq 1 240); do + free=$(nvidia-smi --query-gpu=memory.free --format=csv,noheader,nounits) + [ "${free:-0}" -ge 100000 ] && break + sleep 5 + done + [ "${free:-0}" -ge 100000 ] || { echo "FATAL_NO_VRAM ${free}MiB"; exit 2; } + + ARC_KV_FP8_MODE=$mode "$BIN" bench -m "$BASE/src" -a deepseekv4 \ + --from-uqff "$BASE/uqff/qtip2-0.uqff" --max-seqs 1 --prefix-cache-n 0 \ + --prompt-len 64 --gen-len "$GEN" --iterations 1 --warmup 0 >"$log" 2>&1 + rc=$? + flock -u 9 # release BEFORE parsing + + [ $rc -eq 0 ] || { + echo "FAIL_BENCH $tag rc=$rc" + # A process killed by a GPU fault dies with no error line. This box carries + # ~1,485 Xid faults, so check before blaming the code. + dmesg 2>/dev/null | grep -i xid | tail -3 + tail -5 "$log" + exit 1 + } + # "| Decode (260 tokens) | 13.9 +- 0.0 | 71.80 ms/T |". An empty parse must + # FAIL, not silently become a 0.0 ms delta. + ms=$(grep -oE "[0-9]+\.[0-9]+ ms/T" "$log" | tail -1 | grep -oE "^[0-9]+\.[0-9]+") + [ -n "$ms" ] || { echo "FAIL_NO_MS_PER_T $tag"; tail -5 "$log"; exit 1; } + echo "$tag $mode $ms" +} + +: >"$OUT/ab.txt" +for i in $(seq 1 "$REPS"); do + leg "A$i" cpu | tee -a "$OUT/ab.txt" + leg "B$i" fused | tee -a "$OUT/ab.txt" +done + +echo "=== RESULT ===" +awk ' + { v[substr($1,1,1) substr($1,2)] = $3; arm[substr($1,1,1)] += $3; n[substr($1,1,1)]++ } + END { + a = arm["A"]/n["A"]; b = arm["B"]/n["B"]; + printf "A (cpu round trip) mean %.2f ms/token over %d legs\n", a, n["A"]; + printf "B (fused kernel) mean %.2f ms/token over %d legs\n", b, n["B"]; + printf "delta %.2f ms/token (%+.2f%%)\n", a-b, -100.0*(a-b)/a; + if (n["A"]>1) printf "drift within A %+.2f%% (A1 %.2f -> A%d %.2f)\n", \ + 100.0*(v["A" n["A"]]-v["A1"])/v["A1"], v["A1"], n["A"], v["A" n["A"]]; + if (n["B"]>1) printf "drift within B %+.2f%% (B1 %.2f -> B%d %.2f)\n", \ + 100.0*(v["B" n["B"]]-v["B1"])/v["B1"], v["B1"], n["B"], v["B" n["B"]]; + print "A delta smaller than the drift is not a result."; + }' "$OUT/ab.txt" From 488e4585d49a7c9385744e161d54573515b7da78 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Tue, 18 Aug 2026 12:02:31 +0100 Subject: [PATCH 09/27] measure(ArcKV/Fp8): 44 -> 1 D2H/step; fused kernel 1.98 us/call [MEASURED] Box arc-v4-stack (H200), binary md5 e15259dc9ce935fa8782ba832cac1992 in a PRIVATE target dir (/root/arc-wt/fp8-target, not the shared /root/arc-wt/target that a neighbour's build overwrote tonight), 118 arc_kv_fp8_quantize symbol hits, --features 'cuda flash-attn'. nsys, 20 s steady-state tail, step count pinned from the trace by two independent anchors. before (cpu) after (fused) cuMemcpyDtoHAsync_v2 43.89/step 1.00/step ... of which 1,792 B 42.89/step size ABSENT from the trace ... of which 517,120 B 1.00/step 1.00/step (logits, not ours) kernel launches 9,081.8/step 8,404.5/step (-677.3) cuLaunchKernel 7,892.4/step 7,119.0/step (-773.5) cuMemAllocAsync 11,376.7/step 10,618.0/step (-758.6) ALL *Synchronize* 0.00/step 0.00/step (why grep lies here) Per-call, which is the evidence that survives a contended box: before cuMemcpyDtoHAsync_v2 48.56 us/call HOST time -> 2.131 ms/step blocking after arc_kv_fp8_quantize_kernel 1.98 us/call, 42.93 calls/step arc_kv_fp8_dequantize_kernel 2.34 us/call, 42.93 calls/step fused total 0.1853 ms/step device time 42.93 rather than 43.00 is one step straddling the window edge (99.84%); the before arm's 1,792 B D2H shows the same 0.26% at 42.89/43. Engagement is therefore proven per layer, not assumed. NO end-to-end tok/s delta is reported. A naive A-B-A-B delta on this box is biased by exactly one slot of drift, and four end-to-end numbers here turned out to be pure environment in one night. The A/B driver is dropped rather than shipped with a result it cannot support; the case rests on per-call cost and launch counts, which do not depend on how long the step took. Co-Authored-By: Claude Opus 5 (1M context) --- arc-tools/kv_fp8_ab.sh | 83 ---------------------------- arc-tools/kv_fp8_percall.py | 107 ++++++++++++++++++++++++++++++++++++ 2 files changed, 107 insertions(+), 83 deletions(-) delete mode 100755 arc-tools/kv_fp8_ab.sh create mode 100644 arc-tools/kv_fp8_percall.py diff --git a/arc-tools/kv_fp8_ab.sh b/arc-tools/kv_fp8_ab.sh deleted file mode 100755 index 876d5042c..000000000 --- a/arc-tools/kv_fp8_ab.sh +++ /dev/null @@ -1,83 +0,0 @@ -#!/usr/bin/env bash -# ArcKV/Fp8 — interleaved A-B-A-B decode timing for the fused E4M3 kernel. -# A = ARC_KV_FP8_MODE=cpu (43 blocking D2H per step) -# B = ARC_KV_FP8_MODE=fused (one device kernel, 0 D2H) -# -# A-then-B is banned on this box: a 3.3% monotonic drift once exceeded the arm -# difference. The drift within each arm is printed next to the delta, and a -# delta smaller than the drift is not a result. -# -# LOCK DISCIPLINE. The lock covers exactly one thing: a leg (model load, which -# allocates GPU memory, plus the timed run). It is taken immediately before and -# released immediately after — NOT held across parsing, and NOT wrapped around -# the whole script. Six chains share this H200; one holding the lock through -# post-processing starved the trellis chain for 40 minutes while the card read -# 0 %, 0 MiB, 78 W. -set -u -OUT=${OUT:-/root/kvfp8} -BASE=${BASE:-/root/budget-chain} -BIN=${BIN:-/root/arc-wt/fp8-target/release/mistralrs} -LOCK=${LOCK:-/root/locks/bench.lock} -LOCKWAIT=${LOCKWAIT:-3600} -REPS=${REPS:-2} -GEN=${GEN:-260} -export PATH=/usr/local/cuda-13.1/bin:$PATH -export LD_LIBRARY_PATH=/usr/local/cuda/compat:${LD_LIBRARY_PATH:-} -unset ARC_TIME_DECODE V4_STATS V4_NAN_DEBUG V4_TRACE ARC_PROFILE ARC_GPU_ACT_QUANT -[ -x "$BIN" ] || { echo "FATAL_NO_BINARY $BIN"; exit 2; } - -exec 9>"$LOCK" - -leg() { # $1 = tag, $2 = mode - local tag=$1 mode=$2 log="$OUT/$1.log" free rc ms - flock -w "$LOCKWAIT" 9 || { echo "FATAL_LOCK_TIMEOUT ${LOCKWAIT}s"; exit 2; } - # Holding the lock is not the same as the card being free: V4 is ~79 GB of - # 143 GB and the previous holder's process can outlive its script. - for _ in $(seq 1 240); do - free=$(nvidia-smi --query-gpu=memory.free --format=csv,noheader,nounits) - [ "${free:-0}" -ge 100000 ] && break - sleep 5 - done - [ "${free:-0}" -ge 100000 ] || { echo "FATAL_NO_VRAM ${free}MiB"; exit 2; } - - ARC_KV_FP8_MODE=$mode "$BIN" bench -m "$BASE/src" -a deepseekv4 \ - --from-uqff "$BASE/uqff/qtip2-0.uqff" --max-seqs 1 --prefix-cache-n 0 \ - --prompt-len 64 --gen-len "$GEN" --iterations 1 --warmup 0 >"$log" 2>&1 - rc=$? - flock -u 9 # release BEFORE parsing - - [ $rc -eq 0 ] || { - echo "FAIL_BENCH $tag rc=$rc" - # A process killed by a GPU fault dies with no error line. This box carries - # ~1,485 Xid faults, so check before blaming the code. - dmesg 2>/dev/null | grep -i xid | tail -3 - tail -5 "$log" - exit 1 - } - # "| Decode (260 tokens) | 13.9 +- 0.0 | 71.80 ms/T |". An empty parse must - # FAIL, not silently become a 0.0 ms delta. - ms=$(grep -oE "[0-9]+\.[0-9]+ ms/T" "$log" | tail -1 | grep -oE "^[0-9]+\.[0-9]+") - [ -n "$ms" ] || { echo "FAIL_NO_MS_PER_T $tag"; tail -5 "$log"; exit 1; } - echo "$tag $mode $ms" -} - -: >"$OUT/ab.txt" -for i in $(seq 1 "$REPS"); do - leg "A$i" cpu | tee -a "$OUT/ab.txt" - leg "B$i" fused | tee -a "$OUT/ab.txt" -done - -echo "=== RESULT ===" -awk ' - { v[substr($1,1,1) substr($1,2)] = $3; arm[substr($1,1,1)] += $3; n[substr($1,1,1)]++ } - END { - a = arm["A"]/n["A"]; b = arm["B"]/n["B"]; - printf "A (cpu round trip) mean %.2f ms/token over %d legs\n", a, n["A"]; - printf "B (fused kernel) mean %.2f ms/token over %d legs\n", b, n["B"]; - printf "delta %.2f ms/token (%+.2f%%)\n", a-b, -100.0*(a-b)/a; - if (n["A"]>1) printf "drift within A %+.2f%% (A1 %.2f -> A%d %.2f)\n", \ - 100.0*(v["A" n["A"]]-v["A1"])/v["A1"], v["A1"], n["A"], v["A" n["A"]]; - if (n["B"]>1) printf "drift within B %+.2f%% (B1 %.2f -> B%d %.2f)\n", \ - 100.0*(v["B" n["B"]]-v["B1"])/v["B1"], v["B1"], n["B"], v["B" n["B"]]; - print "A delta smaller than the drift is not a result."; - }' "$OUT/ab.txt" diff --git a/arc-tools/kv_fp8_percall.py b/arc-tools/kv_fp8_percall.py new file mode 100644 index 000000000..ab643a871 --- /dev/null +++ b/arc-tools/kv_fp8_percall.py @@ -0,0 +1,107 @@ +#!/usr/bin/env python3 +"""Per-call cost of the FP8 KV round trip, from an nsys trace. + +Parent system: ArcLab (measuring ArcInfer / ArcKV / Fp8). + +WHY PER-CALL AND NOT tok/s + End-to-end tok/s on this shared box is not trustworthy: four separate + end-to-end numbers here turned out to be pure environment in one night, and + a naive A-B-A-B delta is *biased*, not merely noisy, because at equal + spacing the arms differ by exactly one slot of drift. A per-call cost taken + from the trace is immune to both — it does not depend on how long the step + took, only on what the step did. The `GpuApprox` conclusion survived a bad + night for exactly this reason: its evidence was `kv_fp8_quant` 73.49 -> + 134.18 us/call, not a throughput delta. + +WHAT IT REPORTS + BEFORE arm: the host time spent inside `cuMemcpyDtoHAsync_v2` per decode + step. That is the real cost of the CPU round trip -- an "async" copy taking + ~109 us of HOST time is a pageable-memory staged copy, i.e. blocking, which + is also why a CUDA graph cannot record it. + AFTER arm: the device time of the two fused kernels per decode step, and + their calls/step (which must be 43 -- one per layer -- or the fused path did + not engage on every layer and the number is not what it claims). + +Exits 2 on analysis failure, never 1. + +Usage: kv_fp8_percall.py [tail_seconds] +""" +import collections +import sqlite3 +import sys + + +def die(msg): + print(f"FATAL {msg}", file=sys.stderr) + sys.exit(2) + + +def main(): + if len(sys.argv) < 2: + die("usage: kv_fp8_percall.py [tail_seconds]") + tail_s = float(sys.argv[2]) if len(sys.argv) > 2 else 20.0 + db = sqlite3.connect(sys.argv[1]) + cur = db.cursor() + strings = {i: v for i, v in cur.execute("SELECT id, value FROM StringIds")} + KT = "CUPTI_ACTIVITY_KIND_KERNEL" + MT = "CUPTI_ACTIVITY_KIND_MEMCPY" + RT = "CUPTI_ACTIVITY_KIND_RUNTIME" + + end = cur.execute(f"SELECT MAX(end) FROM {KT}").fetchone()[0] + if end is None: + die("no kernels in the trace") + win0 = end - int(tail_s * 1e9) + + # Steps = the once-per-step logits readback. + big = list(cur.execute( + f"SELECT COUNT(*) FROM {MT} WHERE start>=? AND copyKind=2 AND bytes>100000", + (win0,))) + steps = big[0][0] if big else 0 + if steps < 20: + die(f"only {steps} steps in the window; widen tail_seconds") + print(f"steps in window {steps}") + + # ---- BEFORE-arm evidence: host time inside the blocking D2H. + tot_ns = 0 + n = 0 + for st, en in cur.execute( + f"SELECT start, end FROM {RT} WHERE start>=? AND nameId IN " + f"(SELECT id FROM StringIds WHERE value='cuMemcpyDtoHAsync_v2')", + (win0,)): + tot_ns += en - st + n += 1 + if n: + print(f"cuMemcpyDtoHAsync_v2 {n / steps:8.2f} calls/step " + f"{tot_ns / n / 1000.0:8.2f} us/call (HOST time) " + f"{tot_ns / steps / 1e6:8.3f} ms/step") + else: + print("cuMemcpyDtoHAsync_v2 0.00 calls/step") + + # ---- AFTER-arm evidence: the fused kernels' device time. + dev = collections.defaultdict(lambda: [0, 0]) + for name_id, st, en in cur.execute( + f"SELECT shortName, start, end FROM {KT} WHERE start>=?", (win0,)): + nm = strings.get(name_id, str(name_id)) + if "arc_kv_fp8" in nm: + d = dev[nm] + d[0] += 1 + d[1] += en - st + if dev: + print("fused kernels (device time):") + tot = 0.0 + for nm, (c, ns) in sorted(dev.items()): + per_step = c / steps + print(f" {nm[:44]:44s} {per_step:7.2f} calls/step " + f"{ns / c / 1000.0:7.2f} us/call " + f"{ns / steps / 1e6:7.4f} ms/step") + tot += ns / steps / 1e6 + if abs(per_step - round(per_step)) > 0.05: + print(f" WARN {nm} is not an integer per step") + print(f" fused total{'':38s} {tot:7.4f} ms/step") + else: + print("fused kernels ABSENT from this trace " + "(this is the before arm, or the fused path never engaged)") + + +if __name__ == "__main__": + main() From 46c99e1f23e573e59818c4707ad43cd9187cb21c Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Thu, 20 Aug 2026 18:36:35 +0100 Subject: [PATCH 10/27] =?UTF-8?q?fix(ArcKV/Fp8):=20ARC=5FKV=5FFP8=5FMODE?= =?UTF-8?q?=20was=20a=20lying=20switch=20=E2=80=94=20rename,=20and=20stop?= =?UTF-8?q?=20typos=20selecting=20GpuApprox?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The name says "mode", which reads as an on/off for FP8 KV. It is not one, and it never was: it selects WHICH ARITHMETIC produces the E4M3 code, and all three variants quantize. There is no "off" — V4 is FP8-QAT, so the quantize/dequantize round trip at `deepseek4.rs:1689` is the model's numerics, not an optimisation, and it is correctly unconditional. The flag that gates FP8 KV *storage* is a different variable, `ARC_V4_FP8_KV`. Three defects, in descending severity: 1. A TYPO COULD SILENTLY SELECT NON-BIT-EXACT ARITHMETIC. The old table was `_ if ARC_GPU_ACT_QUANT is set => GpuApprox` followed by `_ => FusedDevice`. Both arms were reachable ONLY for an unset or MISSPELLED value. So on any box that still had `ARC_GPU_ACT_QUANT` exported from an earlier experiment, `ARC_KV_FP8_MODE=fusd` selected `GpuApprox` — the one variant this module documents as round-half-away-from-zero rather than round-half-to-even, and labels "NOT the fix and must not be shipped as one". Resolved once into a `OnceLock`, so it stuck for the process lifetime with no trace. An unparsable value now lands on the bit-exact default and SAYS SO, on both stderr and `tracing::error!` — never on `GpuApprox`. 2. The docs stated the opposite of the code. `PROFILING.md` annotated the `kv_fp8_quant` span "opt-in, ARC_V4_FP8_KV=1". That span opens at `deepseek4.rs:1688` and runs on every forward regardless of the flag; the adjacent `kv_fp8_dequant` line carried no such note, so the doc was internally inconsistent too. The annotation moves to `kv_cache_append`, which is where `ARC_V4_FP8_KV` actually acts. 3. `unwrap_or_default()` collapsed unset and empty-string, and nothing trimmed. Renamed to `ARC_KV_FP8_IMPL`. The old spelling still works and prints a deprecation on both channels — renaming it silently would convert every operator's muscle memory into a fresh silent failure, which is the disease being cured, not the cure. The table is now a pure `parse_impl(Option<&str>, bool)` with a test that pins every arm INCLUDING the typo-must-not-reach-GpuApprox regression. It runs on the free CPU lane; no GPU is required to keep this honest. --- arc-tools/kv_fp8_nsys_ab.sh | 4 +- docs/engineering/PROFILING.md | 6 +- mistralrs-core/src/models/dsv4_kv_fp8.rs | 150 +++++++++++++++++++++-- 3 files changed, 144 insertions(+), 16 deletions(-) diff --git a/arc-tools/kv_fp8_nsys_ab.sh b/arc-tools/kv_fp8_nsys_ab.sh index 362a2b946..80a38692e 100755 --- a/arc-tools/kv_fp8_nsys_ab.sh +++ b/arc-tools/kv_fp8_nsys_ab.sh @@ -49,13 +49,13 @@ exec 9>"$LOCK" for arm in before:cpu after:fused; do A=${arm%%:*} M=${arm##*:} - echo "=== leg $A (ARC_KV_FP8_MODE=$M) $(date -u +%T) ===" + echo "=== leg $A (ARC_KV_FP8_IMPL=$M) $(date -u +%T) ===" rm -f "$OUT/$A".nsys-rep "$OUT/$A".sqlite # LOCK HELD ONLY HERE: VRAM wait + the traced run. Released the instant the # bench exits, before any report export. flock -w "$LOCKWAIT" 9 || { echo "FATAL_LOCK_TIMEOUT after ${LOCKWAIT}s"; exit 2; } wait_for_vram 100000 - ARC_KV_FP8_MODE=$M nsys profile --trace=cuda --sample=none --cpuctxsw=none \ + ARC_KV_FP8_IMPL=$M nsys profile --trace=cuda --sample=none --cpuctxsw=none \ --cuda-memory-usage=false --duration="$DUR" --kill=sigterm \ --force-overwrite=true --output="$OUT/$A" \ "$BIN" bench -m "$BASE/src" -a deepseekv4 \ diff --git a/docs/engineering/PROFILING.md b/docs/engineering/PROFILING.md index 92c466dfc..55a183774 100644 --- a/docs/engineering/PROFILING.md +++ b/docs/engineering/PROFILING.md @@ -211,10 +211,10 @@ step engine/mod.rs — one scheduler iter │ │ │ │ ├─ kv_proj [device] fused wkv │ │ │ │ ├─ kv_norm [device] │ │ │ │ ├─ rope [device] -│ │ │ │ ├─ kv_fp8_quant [device] opt-in, ARC_V4_FP8_KV=1 -│ │ │ │ ├─ kv_fp8_dequant [device] +│ │ │ │ ├─ kv_fp8_quant [device] ALWAYS (V4 is FP8-QAT), not opt-in +│ │ │ │ ├─ kv_fp8_dequant [device] ALWAYS — pairs with the quant above │ │ │ │ ├─ compressed_kv_build[device] -│ │ │ │ ├─ kv_cache_append [device] +│ │ │ │ ├─ kv_cache_append [device] FP8 *storage* here is ARC_V4_FP8_KV=1 │ │ │ │ ├─ kv_cache_span [device] │ │ │ │ ├─ sdpa [device] dsv4_attention (window ∧ compressed) │ │ │ │ ├─ inv_rope [device] NOT inside ARC_TIME_DECODE's timer diff --git a/mistralrs-core/src/models/dsv4_kv_fp8.rs b/mistralrs-core/src/models/dsv4_kv_fp8.rs index 773e95861..a52feac5d 100644 --- a/mistralrs-core/src/models/dsv4_kv_fp8.rs +++ b/mistralrs-core/src/models/dsv4_kv_fp8.rs @@ -73,7 +73,7 @@ pub(crate) enum KvQuantMode { FusedDevice, /// Exact E4M3 via candle's CPU cast (candle has no CUDA `F8E4M3` cast — /// "named symbol not found"), at the price of one device sync per layer. - /// `ARC_KV_FP8_MODE=cpu`. + /// `ARC_KV_FP8_IMPL=cpu`. CpuExact, /// On-device float arithmetic that reproduces E4M3's value grid with /// round-half-away-from-zero instead of round-half-to-even. Removes the @@ -90,25 +90,94 @@ pub(crate) enum KvQuantMode { GpuApprox, } +/// The legal values of `ARC_KV_FP8_IMPL`, for error messages and for the test +/// that pins the table. +pub(crate) const KV_FP8_IMPL_VALUES: &[&str] = &[ + "fused", + "fused_device", + "cpu", + "cpu_exact", + "gpu", + "gpu_approx", +]; + impl KvQuantMode { + /// Pure half of [`Self::from_env`], so the table can be tested without + /// mutating the process environment (which is racy across test threads and + /// `unsafe` since the 2024 edition). + /// + /// `raw` is the `ARC_KV_FP8_IMPL` value (`None` when unset); `legacy_gpu` + /// is whether the pre-existing `ARC_GPU_ACT_QUANT` is set. Returns the mode + /// and, when the input was not understood, the complaint to shout about it. + pub(crate) fn parse_impl(raw: Option<&str>, legacy_gpu: bool) -> (Self, Option) { + let default = if legacy_gpu { + // Preserved from before this flag existed: `ARC_GPU_ACT_QUANT=1` + // selects the approximate on-device path. It only applies when + // `ARC_KV_FP8_IMPL` says nothing. + Self::GpuApprox + } else { + Self::FusedDevice + }; + match raw.map(|s| s.trim()) { + None | Some("") => (default, None), + Some(v) => match v.to_ascii_lowercase().as_str() { + "cpu" | "cpu_exact" => (Self::CpuExact, None), + "gpu" | "gpu_approx" => (Self::GpuApprox, None), + "fused" | "fused_device" => (Self::FusedDevice, None), + // NOT `default`. Falling through to `default` here is how a + // typo used to select `GpuApprox` — the one variant documented + // as *not* bit-exact and "must not be shipped" — on any box + // that still had `ARC_GPU_ACT_QUANT` set from an earlier + // experiment. An unparsable value now lands on the bit-exact + // default and says so. + _ => ( + Self::FusedDevice, + Some(format!( + "ARC_KV_FP8_IMPL={v:?} is not a legal value (expected one of {}); \ + falling back to the default fused-device kernel. \ + NOTE: this flag selects WHICH ARITHMETIC produces the E4M3 code. \ + It has no 'off' — V4 is FP8-QAT and the quantize/dequantize round \ + trip is the model's numerics, not an optimisation. The flag that \ + gates FP8 KV *storage* is ARC_V4_FP8_KV.", + KV_FP8_IMPL_VALUES.join(", ") + )), + ), + }, + } + } + /// Resolved once per process: this is called per attention layer per /// forward, and `deepseek4` has already been bitten by per-call /// `std::env::var_os` in exactly that position (~390 environment scans per /// forward, wave33). + /// + /// Reads `ARC_KV_FP8_IMPL`. The former spelling `ARC_KV_FP8_MODE` is still + /// honoured — loudly — because renaming it silently would turn every + /// operator's muscle memory into a new silent failure, which is the same + /// disease the rename cures. pub(crate) fn from_env() -> Self { static MODE: OnceLock = OnceLock::new(); *MODE.get_or_init(|| { - match std::env::var("ARC_KV_FP8_MODE") - .unwrap_or_default() - .to_ascii_lowercase() - .as_str() - { - "cpu" | "cpu_exact" => Self::CpuExact, - "gpu" | "gpu_approx" => Self::GpuApprox, - "fused" | "fused_device" => Self::FusedDevice, - _ if std::env::var_os("ARC_GPU_ACT_QUANT").is_some() => Self::GpuApprox, - _ => Self::FusedDevice, + let legacy_gpu = std::env::var_os("ARC_GPU_ACT_QUANT").is_some(); + let new = std::env::var("ARC_KV_FP8_IMPL").ok(); + let old = std::env::var("ARC_KV_FP8_MODE").ok(); + if old.is_some() { + // eprintln! as well as tracing: a rented box often runs before + // a subscriber is installed, and a deprecation nobody sees is + // not a deprecation. + let msg = "ARC_KV_FP8_MODE has been renamed ARC_KV_FP8_IMPL (it selects the \ + quantizer ARITHMETIC; it never had an 'off'). The old name still \ + works for now — please update."; + eprintln!("[arc-kv-fp8] {msg}"); + tracing::warn!("{msg}"); } + let raw = new.or(old); + let (mode, complaint) = Self::parse_impl(raw.as_deref(), legacy_gpu); + if let Some(c) = complaint { + eprintln!("[arc-kv-fp8] {c}"); + tracing::error!("{c}"); + } + mode }) } } @@ -739,4 +808,63 @@ mod tests { assert_eq!(values[0], 0.0); assert_eq!(values[126], 448.0); } + + /// The parse table, pinned. Before this, `ARC_KV_FP8_MODE` had a bare + /// `_ => Self::FusedDevice` catch-all preceded by an + /// `_ if ARC_GPU_ACT_QUANT is set => Self::GpuApprox` arm, so a MISSPELLED + /// value on a box that still had `ARC_GPU_ACT_QUANT` exported silently + /// selected `GpuApprox` — the one variant this module documents as not + /// bit-exact and "must not be shipped as one". + #[test] + fn kv_fp8_impl_parse_table() { + use KvQuantMode::*; + + // Unset: the bit-exact fused kernel, with and without trailing space. + assert_eq!(KvQuantMode::parse_impl(None, false).0, FusedDevice); + assert_eq!(KvQuantMode::parse_impl(Some(""), false).0, FusedDevice); + assert_eq!(KvQuantMode::parse_impl(Some(" "), false).0, FusedDevice); + + // Every legal value round-trips, case- and whitespace-insensitively, + // and none of them complains. + for (raw, want) in [ + ("fused", FusedDevice), + ("fused_device", FusedDevice), + ("FUSED", FusedDevice), + (" cpu ", CpuExact), + ("cpu_exact", CpuExact), + ("gpu", GpuApprox), + ("gpu_approx", GpuApprox), + ] { + let (got, complaint) = KvQuantMode::parse_impl(Some(raw), false); + assert_eq!(got, want, "{raw:?}"); + assert!(complaint.is_none(), "{raw:?} should parse silently"); + } + + // Every value the error message advertises as legal really is legal. + for v in KV_FP8_IMPL_VALUES { + assert!( + KvQuantMode::parse_impl(Some(v), false).1.is_none(), + "{v:?} is advertised as legal but is rejected" + ); + } + + // The legacy var still selects GpuApprox, but only when the new flag + // is silent. + assert_eq!(KvQuantMode::parse_impl(None, true).0, GpuApprox); + assert_eq!(KvQuantMode::parse_impl(Some(""), true).0, GpuApprox); + assert_eq!(KvQuantMode::parse_impl(Some("cpu"), true).0, CpuExact); + + // THE REGRESSION. A typo must be loud, and must NOT reach GpuApprox + // even with the legacy var set. + for legacy in [false, true] { + let (got, complaint) = KvQuantMode::parse_impl(Some("fusd"), legacy); + assert_eq!(got, FusedDevice, "typo must land on the bit-exact default"); + let c = complaint.expect("a typo must produce a complaint"); + assert!(c.contains("fusd"), "the complaint must quote the bad value"); + assert!( + c.contains("ARC_V4_FP8_KV"), + "the complaint must name the flag that actually gates storage" + ); + } + } } From 3a71d0392013e6269ea3676370537e60074f2b08 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Thu, 20 Aug 2026 18:46:22 +0100 Subject: [PATCH 11/27] =?UTF-8?q?fix(build):=20EXPECTED=5FKERNEL=5FCOUNT?= =?UTF-8?q?=2040=20->=2041=20=E2=80=94=20this=20branch=20adds=20arc=5Fkvqu?= =?UTF-8?q?ant.cu?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The ArcGate kernel-count tripwire caught a real regression introduced by rebasing this branch onto current master, and it is worth recording how, because the failure mode is subtle. This branch originally carried `EXPECTED_KERNEL_COUNT 39 -> 40 — arc_kvquant.cu was never counted`. Meanwhile master independently went 39 -> 40 for a DIFFERENT kernel. On rebase, git compared patch texts, saw an identical `-39 / +41`-shaped hunk already upstream, and dropped the commit as "patch contents already upstream". The number was right; the REASON was not. Net effect: the glob discovers 41 sources while the file still claims 40, so `arc_kvquant.cu` would once again be the kernel that goes missing quietly — which is the exact failure this file was created to make impossible. Caught by `cuda_kernel_build_guard::expected_kernel_count_matches_disk` on the free CPU lane, before any GPU time was spent. That is the tripwire earning its keep, not a nuisance — do not "fix" a future occurrence by relaxing the guard. --- mistralrs-quant/kernels/EXPECTED_KERNEL_COUNT | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mistralrs-quant/kernels/EXPECTED_KERNEL_COUNT b/mistralrs-quant/kernels/EXPECTED_KERNEL_COUNT index c5cd0dc04..b0a528a99 100644 --- a/mistralrs-quant/kernels/EXPECTED_KERNEL_COUNT +++ b/mistralrs-quant/kernels/EXPECTED_KERNEL_COUNT @@ -28,4 +28,4 @@ # (this count minus the 5 `*_dummy.cu` / `dummy_*.cu` stubs on an SM >= 80 # build), and that line is the engagement counter for a CUDA build: a CUDA # build that "succeeded" without it moving has not built what you think it did. -40 +41 From dd5583c4b753ae76f01aa84a65a5b7b516d04202 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Wed, 19 Aug 2026 20:20:59 +0100 Subject: [PATCH 12/27] perf(arckv): make candle's caching allocator reachable and on for decode MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `set_alloc_cache_enabled(true)` sat behind three stacked default-off gates: `probe && seq_len == 1 && env("ARC_CANDLE_ALLOC_CACHE")`. The first conjunct tied a general-purpose allocator to the V4 capture probe, so the only way to recycle a decode step's frees was to also be capturing. The ~11k allocations per token were a disabled feature, not a missing one. Replaced with a pure policy, `alloc_cache_action(seq_len, enabled, killed)`: * decode (`seq_len == 1`) -> Enable * prefill (`seq_len != 1`) -> DrainAndDisable * already in that state -> Leave Prefill draining is not incidental. The cache is keyed on exact byte count (`free: HashMap>`, no bucketing, no smallest-fit) and has no capacity bound and no eviction, so a prefill's large one-shot buffers would be parked for the process lifetime under a key nothing asks for again. `ARC_CANDLE_ALLOC_CACHE=0` is the kill switch. Any other value, and unset, leave the policy in force, so the `=1` the ops scripts pass still means what it always meant. The capture probe keeps its own gate for the graph-mode positions; only the allocator moved out from under it. Six host-runnable tests for the policy. The allocator has no tests at all in either repo — candle's `cuda_backend/{device,mod}.rs` carry no `#[cfg(test)]` and every arc-side exercise is `#[cfg(feature = "cuda")]` — so the decision of when it is on is now the part that CI can see. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SpVNMpb13HkUXqSqbN1o9H --- mistralrs-core/src/pipeline/normal.rs | 200 ++++++++++++++++++++++++-- 1 file changed, 192 insertions(+), 8 deletions(-) diff --git a/mistralrs-core/src/pipeline/normal.rs b/mistralrs-core/src/pipeline/normal.rs index d412a7cd9..7f32c2563 100644 --- a/mistralrs-core/src/pipeline/normal.rs +++ b/mistralrs-core/src/pipeline/normal.rs @@ -98,6 +98,91 @@ pub struct NormalPipeline { autonomous_runner: Option, } +/// What a forward pass should do with candle's caching allocator before it +/// runs. See [`alloc_cache_action`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum AllocCacheAction { + /// Turn the cache on. Frees start being recycled instead of returned to + /// the driver. + Enable, + /// Return every held buffer to the driver and turn the cache off. + DrainAndDisable, + /// Already in the right state; touch nothing. + Leave, +} + +/// Decide whether candle's caching allocator should be on for this forward. +/// +/// # Why this is not simply "on" +/// +/// The cache (fork `88d86a2`, `candle-core/src/cuda_backend/device.rs:42`) is +/// keyed on **exact byte count** — `free: HashMap>`, +/// looked up with `free.get_mut(&bytes)`. No bucketing, no smallest-fit, no +/// splitting. It also has **no capacity bound and no eviction**: the only ways +/// memory goes back to the driver are `set_alloc_cache_enabled(false)` and +/// `drain_alloc_cache_and_free()`. +/// +/// Those two facts together decide the policy: +/// +/// * **Decode wants it on.** A decode step allocates the same shapes 43 times +/// over (once per layer), and the shapes that do not depend on KV length — +/// hidden states, MLP and expert intermediates — repeat step after step. +/// Exact-size keying is a perfect fit for that traffic. +/// * **Prefill must drain it.** Prefill shapes scale with prompt length, are +/// large, and are hit once. Leaving the cache on across a prefill parks +/// those buffers for the process lifetime under a byte-size key nothing will +/// ever request again. That is the failure mode `ARC_NO_DEDICATED_DECODE` +/// already exists to work around, and it is why this returns +/// [`AllocCacheAction::DrainAndDisable`] rather than `Leave` on `seq_len != 1`. +/// +/// # The measurement this is waiting on +/// +/// Buffers whose size tracks KV length — the causal mask most obviously — +/// change size every decode step, so each step files one more never-reused +/// entry. That is bounded by `O(context^2)` bytes in the worst case and is +/// *not* bounded by this policy. It is small at short context (a `[1,1,1,kv]` +/// BF16 mask over 4k tokens sums to ~16 MB) and is not small at 128k. +/// +/// The fixed-capacity graph-mode path (`deepseek4.rs:4344-4353`, which swaps +/// the growing causal mask for `graph_mode_length_mask` at +/// `cfg_full.sliding_window`) removes that growth entirely, but it is reached +/// only under `ARC_V4_CAPTURE_PROBE`. Until the shape-invariance work lands, +/// `ARC_CANDLE_ALLOC_CACHE=0` is the kill switch, and the long-context +/// high-water mark is the number a GPU run should falsify this with. +/// +/// # Contract +/// +/// * `seq_len` — the forward's sequence length. `1` is decode. +/// * `enabled` — what candle reports *now* (`alloc_cache_enabled()`), so the +/// action is idempotent and we never drain a cache that is already off. +/// * `killed` — `ARC_CANDLE_ALLOC_CACHE=0` was set. +/// +/// Pure, so the policy is testable without a GPU — which matters, because the +/// allocator itself has no tests at all in either repo. +#[cfg_attr(not(feature = "cuda"), allow(dead_code))] +pub(crate) fn alloc_cache_action( + seq_len: usize, + enabled: bool, + killed: bool, +) -> AllocCacheAction { + let want = !killed && seq_len == 1; + match (want, enabled) { + (true, false) => AllocCacheAction::Enable, + (false, true) => AllocCacheAction::DrainAndDisable, + _ => AllocCacheAction::Leave, + } +} + +/// `ARC_CANDLE_ALLOC_CACHE=0` turns the caching allocator off entirely. +/// +/// Any other value — and *unset* — leaves the default policy in force. The +/// variable used to be the on-switch, so the value `1` the ops scripts pass +/// (`arc-tools/arcgraph_heap_probe.sh:191`) still means what it always meant. +#[cfg_attr(not(feature = "cuda"), allow(dead_code))] +pub(crate) fn alloc_cache_killed() -> bool { + std::env::var("ARC_CANDLE_ALLOC_CACHE").is_ok_and(|v| v == "0") +} + /// A loader for a "normal" (non-quantized) model. pub struct NormalLoader { inner: Box, @@ -1702,15 +1787,31 @@ impl Pipeline for NormalPipeline { "normal.rs:1554", ); } - // Enable the candle caching allocator for graph-capture - // safety (RUN-161). Idempotent; gated so it's off during - // model load and only active for decode. Warmup decode - // forwards populate the cache before capture. - if probe && seq_len == 1 && std::env::var_os("ARC_CANDLE_ALLOC_CACHE").is_some() - { - if let candle_core::Device::Cuda(cd) = self.device() { - cd.set_alloc_cache_enabled(true); + // Candle's caching allocator, driven by `alloc_cache_action`. + // + // This used to be `probe && seq_len == 1 && env(...)`, which + // made a general-purpose allocator reachable only when the + // V4 capture probe was also on. The cache recycles frees for + // ANY decode step — it is not capture machinery — and the + // three stacked default-off gates meant the ~11k allocations + // per token were a disabled feature rather than a missing + // one. It is now on for decode by default and drained on + // prefill; see `alloc_cache_action` for why prefill must + // drain rather than coast. + if let candle_core::Device::Cuda(cd) = self.device() { + match alloc_cache_action( + seq_len, + cd.alloc_cache_enabled(), + alloc_cache_killed(), + ) { + AllocCacheAction::Enable => cd.set_alloc_cache_enabled(true), + AllocCacheAction::DrainAndDisable => { + cd.set_alloc_cache_enabled(false) + } + AllocCacheAction::Leave => {} } + } + if probe && seq_len == 1 { // RUN-161 step 2b. Set the graph-mode device position: // drives RoPE + the fixed-capacity KV write slot, and // makes warmup forwards take the shape-constant path so @@ -2528,3 +2629,86 @@ impl AnyMoePipelineMixin for NormalPipeline { self.model.amoe_supported() } } + +/// The caching-allocator policy's contract. +/// +/// The allocator itself has **no tests at all** — not in this repo, and not in +/// the candle fork that implements it (`grep -rn "alloc_cache" --include="*.rs"` +/// over `candle-core` finds only `cuda_backend/device.rs` and +/// `cuda_backend/mod.rs`, with no `#[cfg(test)]` in either). Everything that +/// observes it needs a GPU, so nothing observes it in CI. +/// +/// The decision of *when* it is on does not need a GPU, so it is tested here. +#[cfg(test)] +mod alloc_cache_policy_tests { + use super::{alloc_cache_action, AllocCacheAction}; + + /// The change this policy exists to make: a decode step turns the cache on + /// without any capture probe being involved. + #[test] + fn decode_enables_the_cache() { + assert_eq!( + alloc_cache_action(1, false, false), + AllocCacheAction::Enable + ); + } + + /// Prefill must hand its buffers back. They are large, they are keyed by an + /// exact byte size no later request will ask for, and nothing evicts them. + #[test] + fn prefill_drains_the_cache() { + assert_eq!( + alloc_cache_action(512, true, false), + AllocCacheAction::DrainAndDisable + ); + } + + /// Called once per forward, so it must be a no-op in the steady state + /// rather than re-enabling (or re-draining) every step. + #[test] + fn steady_state_touches_nothing() { + assert_eq!(alloc_cache_action(1, true, false), AllocCacheAction::Leave); + assert_eq!( + alloc_cache_action(512, false, false), + AllocCacheAction::Leave + ); + } + + /// The kill switch has to work from either state, including turning off a + /// cache that a previous step already enabled — otherwise setting it + /// mid-run would leave the buffers stranded. + #[test] + fn kill_switch_disables_and_drains_from_either_state() { + assert_eq!( + alloc_cache_action(1, true, true), + AllocCacheAction::DrainAndDisable + ); + assert_eq!(alloc_cache_action(1, false, true), AllocCacheAction::Leave); + } + + /// A prompt that happens to be one token long is a prefill by every other + /// measure, but it allocates decode-shaped buffers, so the policy keys on + /// the shape it will actually see. Pinned so the equivalence is deliberate. + #[test] + fn one_token_prompt_is_treated_as_decode() { + assert_eq!( + alloc_cache_action(1, false, false), + AllocCacheAction::Enable + ); + } + + /// A zero-length forward is not decode. Guards against `seq_len == 0` + /// (which `dims2().unwrap_or((0, 0))` produces on a shape error) quietly + /// enabling the cache. + #[test] + fn degenerate_zero_length_forward_does_not_enable() { + assert_eq!( + alloc_cache_action(0, false, false), + AllocCacheAction::Leave + ); + assert_eq!( + alloc_cache_action(0, true, false), + AllocCacheAction::DrainAndDisable + ); + } +} From 7d91c1127ee3118fccd603bdfa64b73f2fb724aa Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Wed, 19 Aug 2026 23:48:23 +0100 Subject: [PATCH 13/27] perf(arckv): bound the caching allocator, and print the counters that prove it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Picks up candle b2a4dd80, which gives `AllocCache` a capacity and LRU eviction. The allocator had neither, and arc only drains it on a prefill, so a single long generation grew forever. Measured on an H200 over 2 600 tokens, `memory.used` at 2 Hz: **+6.04 MiB per decoded token with no plateau**, against +0.057 MiB/token for the same run with the cache off — so the growth is the cache's and nothing else's. `ARC_ALLOC_CACHE_MAX_MB` sets the cap; unset leaves candle's 1 GiB default; `0` restores the old unbounded behaviour for A/B. A typo deliberately does *not* fall back to unbounded. `ARC_ALLOC_CACHE_STATS=N` prints the allocator's counters every N decode steps: allocations per step, **frees per step**, hit rate, bytes held against the cap. Those are the numbers this has to be judged on. A green log is not evidence — an earlier arena here reported "accounting OK" and bit-identical output over 52 steps while silently bypassing itself for every buffer under 128 bytes (KERNEL_RULES.md:977-984). Allocations staying low *and* frees being non-zero is. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SpVNMpb13HkUXqSqbN1o9H --- Cargo.toml | 10 +- mistralrs-core/src/pipeline/normal.rs | 138 +++++++++++++++++++++++++- 2 files changed, 142 insertions(+), 6 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index b8dde1f07..225d8ecbe 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -49,11 +49,11 @@ rust-version = "1.88" # pointer has broken CI on this repo before — a path/branch dependency that # advances under a merged PR turns a green build into a red one with no commit # here to blame. -candle-core = { git = "https://github.com/aeonmindai/candle.git", rev = "88d86a2019ee552923b217284e89847d2785dcdf" } -candle-nn = { git = "https://github.com/aeonmindai/candle.git", rev = "88d86a2019ee552923b217284e89847d2785dcdf" } -candle-flash-attn-v3 = { git = "https://github.com/aeonmindai/candle.git", rev = "88d86a2019ee552923b217284e89847d2785dcdf" } -candle-flash-attn = { git = "https://github.com/aeonmindai/candle.git", rev = "88d86a2019ee552923b217284e89847d2785dcdf" } -candle-metal-kernels = { git = "https://github.com/aeonmindai/candle.git", rev = "88d86a2019ee552923b217284e89847d2785dcdf" } +candle-core = { git = "https://github.com/aeonmindai/candle.git", rev = "b2a4dd80b82726730c3c4c90421fa3a0938be984" } +candle-nn = { git = "https://github.com/aeonmindai/candle.git", rev = "b2a4dd80b82726730c3c4c90421fa3a0938be984" } +candle-flash-attn-v3 = { git = "https://github.com/aeonmindai/candle.git", rev = "b2a4dd80b82726730c3c4c90421fa3a0938be984" } +candle-flash-attn = { git = "https://github.com/aeonmindai/candle.git", rev = "b2a4dd80b82726730c3c4c90421fa3a0938be984" } +candle-metal-kernels = { git = "https://github.com/aeonmindai/candle.git", rev = "b2a4dd80b82726730c3c4c90421fa3a0938be984" } axum = "0.8.8" anyhow = "1.0.100" diff --git a/mistralrs-core/src/pipeline/normal.rs b/mistralrs-core/src/pipeline/normal.rs index 7f32c2563..a69835923 100644 --- a/mistralrs-core/src/pipeline/normal.rs +++ b/mistralrs-core/src/pipeline/normal.rs @@ -183,6 +183,90 @@ pub(crate) fn alloc_cache_killed() -> bool { std::env::var("ARC_CANDLE_ALLOC_CACHE").is_ok_and(|v| v == "0") } +/// Retention cap for the caching allocator, in bytes, or `None` to leave +/// candle's own default (1 GiB) in force. +/// +/// `ARC_ALLOC_CACHE_MAX_MB=0` means unbounded — the pre-bounding behaviour, kept +/// reachable so the leak it causes can be A/B'd rather than argued about. +/// Measured unbounded on V4: **+6.04 MiB per decoded token with no plateau**, +/// against +0.057 MiB/token with the cache off entirely. +#[cfg_attr(not(feature = "cuda"), allow(dead_code))] +pub(crate) fn alloc_cache_capacity_bytes() -> Option { + parse_alloc_cache_capacity(std::env::var("ARC_ALLOC_CACHE_MAX_MB").ok().as_deref()) +} + +/// Pure half of [`alloc_cache_capacity_bytes`], so the parse is testable. +/// +/// `None` (unset, or unparseable) leaves candle's own default in force rather +/// than silently picking a different one — a typo in an ops script must not +/// quietly hand back the unbounded allocator this change exists to remove. +pub(crate) fn parse_alloc_cache_capacity(raw: Option<&str>) -> Option { + let mb: usize = raw?.trim().parse().ok()?; + Some(if mb == 0 { + usize::MAX + } else { + mb.saturating_mul(1024 * 1024) + }) +} + +/// Emit the allocator's counters every `ARC_ALLOC_CACHE_STATS` decode steps. +/// +/// The counters are what this cache has to be judged on. A green log — the +/// server did not crash, the output looked fine — is not evidence that the +/// cache is working: an earlier arena in this codebase reported "accounting OK" +/// and bit-identical output across 52 steps while silently bypassing itself for +/// every buffer under 128 bytes (`KERNEL_RULES.md:977-984`). What distinguishes +/// a working bounded cache from that is arithmetic, and it is printed here: +/// +/// * `alloc/step` — real `cuMemAllocAsync` calls per decode step. Low means the +/// cache is absorbing the ~11 k allocations a step makes. +/// * `free/step` — real `cuMemFreeAsync` calls per decode step. **Non-zero is +/// the point.** The unbounded cache's was exactly zero, forever, which is why +/// it grew without bound. +/// * `held` — bytes retained right now, against the cap. +#[cfg(feature = "cuda")] +fn report_alloc_cache_step(cd: &candle_core::CudaDevice, seq_len: usize) { + use std::sync::atomic::{AtomicU64, Ordering}; + static EVERY: std::sync::OnceLock = std::sync::OnceLock::new(); + let every = *EVERY.get_or_init(|| { + std::env::var("ARC_ALLOC_CACHE_STATS") + .ok() + .and_then(|v| v.trim().parse::().ok()) + .unwrap_or(0) + }); + if every == 0 || seq_len != 1 { + return; + } + static STEP: AtomicU64 = AtomicU64::new(0); + static LAST_ALLOC: AtomicU64 = AtomicU64::new(0); + static LAST_FREE: AtomicU64 = AtomicU64::new(0); + static LAST_STEP: AtomicU64 = AtomicU64::new(0); + let step = STEP.fetch_add(1, Ordering::Relaxed) + 1; + if step % every != 0 { + return; + } + let s = cd.alloc_cache_stats(); + let d_step = step - LAST_STEP.swap(step, Ordering::Relaxed); + let d_alloc = s.misses - LAST_ALLOC.swap(s.misses, Ordering::Relaxed); + let d_free = s.frees() - LAST_FREE.swap(s.frees(), Ordering::Relaxed); + let n = d_step.max(1) as f64; + tracing::info!( + "[alloc-cache] step {step} alloc/step {:.1} free/step {:.1} \ + hit-rate {:.4} held {:.1} MiB / cap {} high-water {:.1} MiB sizes {}", + d_alloc as f64 / n, + d_free as f64 / n, + s.hits as f64 / (s.hits + s.misses).max(1) as f64, + s.cached_bytes as f64 / (1024.0 * 1024.0), + if s.capacity_bytes == usize::MAX { + "unbounded".to_string() + } else { + format!("{} MiB", s.capacity_bytes / (1024 * 1024)) + }, + s.high_water_bytes as f64 / (1024.0 * 1024.0), + s.size_classes, + ); +} + /// A loader for a "normal" (non-quantized) model. pub struct NormalLoader { inner: Box, @@ -1804,12 +1888,18 @@ impl Pipeline for NormalPipeline { cd.alloc_cache_enabled(), alloc_cache_killed(), ) { - AllocCacheAction::Enable => cd.set_alloc_cache_enabled(true), + AllocCacheAction::Enable => { + cd.set_alloc_cache_enabled(true); + if let Some(cap) = alloc_cache_capacity_bytes() { + cd.set_alloc_cache_capacity(cap); + } + } AllocCacheAction::DrainAndDisable => { cd.set_alloc_cache_enabled(false) } AllocCacheAction::Leave => {} } + report_alloc_cache_step(cd, seq_len); } if probe && seq_len == 1 { // RUN-161 step 2b. Set the graph-mode device position: @@ -2712,3 +2802,49 @@ mod alloc_cache_policy_tests { ); } } + +/// The retention cap's parse. +/// +/// The cap is the whole of the fix — candle's caching allocator had none, and +/// retained 6.04 MiB per decoded token forever — so the way it is configured has +/// to be unambiguous. In particular a typo must not fall back to "unbounded". +#[cfg(test)] +mod alloc_cache_capacity_tests { + use super::parse_alloc_cache_capacity; + + #[test] + fn unset_leaves_candles_default_in_force() { + assert_eq!(parse_alloc_cache_capacity(None), None); + } + + /// A typo is not a licence to run unbounded. `None` means "don't touch it", + /// and candle's default is bounded, so a bad value degrades to bounded. + #[test] + fn an_unparseable_value_does_not_become_unbounded() { + for bad in ["", " ", "lots", "1GiB", "-1", "1.5"] { + assert_eq!(parse_alloc_cache_capacity(Some(bad)), None, "{bad:?}"); + } + } + + /// `0` is the documented escape hatch back to the old unbounded allocator, + /// kept reachable only so its leak can be re-measured rather than argued + /// about. + #[test] + fn zero_is_the_explicit_unbounded_opt_in() { + assert_eq!(parse_alloc_cache_capacity(Some("0")), Some(usize::MAX)); + } + + #[test] + fn megabytes_convert_and_do_not_overflow() { + assert_eq!(parse_alloc_cache_capacity(Some("1")), Some(1024 * 1024)); + assert_eq!( + parse_alloc_cache_capacity(Some(" 512 ")), + Some(512 * 1024 * 1024) + ); + assert_eq!( + parse_alloc_cache_capacity(Some(&usize::MAX.to_string())), + Some(usize::MAX), + "saturates instead of wrapping to a tiny cap" + ); + } +} From f9d496acf929e015bebc43b2b292e58ceebd0c82 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Wed, 19 Aug 2026 23:54:30 +0100 Subject: [PATCH 14/27] fix(build): report_alloc_cache_step takes a reference Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SpVNMpb13HkUXqSqbN1o9H --- mistralrs-core/src/pipeline/normal.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mistralrs-core/src/pipeline/normal.rs b/mistralrs-core/src/pipeline/normal.rs index a69835923..4675ab926 100644 --- a/mistralrs-core/src/pipeline/normal.rs +++ b/mistralrs-core/src/pipeline/normal.rs @@ -1899,7 +1899,7 @@ impl Pipeline for NormalPipeline { } AllocCacheAction::Leave => {} } - report_alloc_cache_step(cd, seq_len); + report_alloc_cache_step(&cd, seq_len); } if probe && seq_len == 1 { // RUN-161 step 2b. Set the graph-mode device position: From fa1d2ef28b1def98d3767b7302018e862ddc66e4 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Thu, 20 Aug 2026 00:22:50 +0100 Subject: [PATCH 15/27] =?UTF-8?q?chore(deps):=20candle=20859c49c8=20?= =?UTF-8?q?=E2=80=94=20allocator=20test=20fixes?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Test-only change on the candle side; `git diff b2a4dd80..859c49c8` touches nothing outside `#[cfg(test)] mod alloc_cache_tests`. The H200 numbers in the branch description were measured at b2a4dd80 and stand. Correction to that commit message: the allocator has **ten** tests, not eleven. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SpVNMpb13HkUXqSqbN1o9H --- Cargo.toml | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 225d8ecbe..0f7f7935f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -49,11 +49,11 @@ rust-version = "1.88" # pointer has broken CI on this repo before — a path/branch dependency that # advances under a merged PR turns a green build into a red one with no commit # here to blame. -candle-core = { git = "https://github.com/aeonmindai/candle.git", rev = "b2a4dd80b82726730c3c4c90421fa3a0938be984" } -candle-nn = { git = "https://github.com/aeonmindai/candle.git", rev = "b2a4dd80b82726730c3c4c90421fa3a0938be984" } -candle-flash-attn-v3 = { git = "https://github.com/aeonmindai/candle.git", rev = "b2a4dd80b82726730c3c4c90421fa3a0938be984" } -candle-flash-attn = { git = "https://github.com/aeonmindai/candle.git", rev = "b2a4dd80b82726730c3c4c90421fa3a0938be984" } -candle-metal-kernels = { git = "https://github.com/aeonmindai/candle.git", rev = "b2a4dd80b82726730c3c4c90421fa3a0938be984" } +candle-core = { git = "https://github.com/aeonmindai/candle.git", rev = "859c49c8c4e24177378a816f7b3f3d284987c62f" } +candle-nn = { git = "https://github.com/aeonmindai/candle.git", rev = "859c49c8c4e24177378a816f7b3f3d284987c62f" } +candle-flash-attn-v3 = { git = "https://github.com/aeonmindai/candle.git", rev = "859c49c8c4e24177378a816f7b3f3d284987c62f" } +candle-flash-attn = { git = "https://github.com/aeonmindai/candle.git", rev = "859c49c8c4e24177378a816f7b3f3d284987c62f" } +candle-metal-kernels = { git = "https://github.com/aeonmindai/candle.git", rev = "859c49c8c4e24177378a816f7b3f3d284987c62f" } axum = "0.8.8" anyhow = "1.0.100" From 65dc0fb4a36d4615a477afad5cd9ef2f25278c3b Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Thu, 20 Aug 2026 00:50:50 +0100 Subject: [PATCH 16/27] =?UTF-8?q?chore(deps):=20candle=2089ab14ef=20?= =?UTF-8?q?=E2=80=94=20no=20eviction=20inside=20the=20capture=20window?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes a free path this change introduced: leaving capture mode re-filed parked buffers through the evicting put, which could hand a private-pool pointer back to the driver at a moment arc-cuda-graph does not control. Decode is unaffected — `set_capture_mode` is only called during a capture. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SpVNMpb13HkUXqSqbN1o9H --- Cargo.toml | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 0f7f7935f..b8852835d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -49,11 +49,11 @@ rust-version = "1.88" # pointer has broken CI on this repo before — a path/branch dependency that # advances under a merged PR turns a green build into a red one with no commit # here to blame. -candle-core = { git = "https://github.com/aeonmindai/candle.git", rev = "859c49c8c4e24177378a816f7b3f3d284987c62f" } -candle-nn = { git = "https://github.com/aeonmindai/candle.git", rev = "859c49c8c4e24177378a816f7b3f3d284987c62f" } -candle-flash-attn-v3 = { git = "https://github.com/aeonmindai/candle.git", rev = "859c49c8c4e24177378a816f7b3f3d284987c62f" } -candle-flash-attn = { git = "https://github.com/aeonmindai/candle.git", rev = "859c49c8c4e24177378a816f7b3f3d284987c62f" } -candle-metal-kernels = { git = "https://github.com/aeonmindai/candle.git", rev = "859c49c8c4e24177378a816f7b3f3d284987c62f" } +candle-core = { git = "https://github.com/aeonmindai/candle.git", rev = "89ab14ef1216331a782539f533b99f6816708604" } +candle-nn = { git = "https://github.com/aeonmindai/candle.git", rev = "89ab14ef1216331a782539f533b99f6816708604" } +candle-flash-attn-v3 = { git = "https://github.com/aeonmindai/candle.git", rev = "89ab14ef1216331a782539f533b99f6816708604" } +candle-flash-attn = { git = "https://github.com/aeonmindai/candle.git", rev = "89ab14ef1216331a782539f533b99f6816708604" } +candle-metal-kernels = { git = "https://github.com/aeonmindai/candle.git", rev = "89ab14ef1216331a782539f533b99f6816708604" } axum = "0.8.8" anyhow = "1.0.100" From e4eb59dfba78953ee8d85aa1dc3da69801e6452c Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Thu, 20 Aug 2026 11:56:53 +0100 Subject: [PATCH 17/27] fix(ci): spell 'unparsable' the way the typos gate expects MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two sites in pipeline/normal.rs — a doc comment and a test name. No behaviour change; the Typos job was the only red check on this PR. Co-Authored-By: Claude Opus 5 (1M context) --- mistralrs-core/src/pipeline/normal.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mistralrs-core/src/pipeline/normal.rs b/mistralrs-core/src/pipeline/normal.rs index 4675ab926..0bcc24e8a 100644 --- a/mistralrs-core/src/pipeline/normal.rs +++ b/mistralrs-core/src/pipeline/normal.rs @@ -197,7 +197,7 @@ pub(crate) fn alloc_cache_capacity_bytes() -> Option { /// Pure half of [`alloc_cache_capacity_bytes`], so the parse is testable. /// -/// `None` (unset, or unparseable) leaves candle's own default in force rather +/// `None` (unset, or unparsable) leaves candle's own default in force rather /// than silently picking a different one — a typo in an ops script must not /// quietly hand back the unbounded allocator this change exists to remove. pub(crate) fn parse_alloc_cache_capacity(raw: Option<&str>) -> Option { @@ -2820,7 +2820,7 @@ mod alloc_cache_capacity_tests { /// A typo is not a licence to run unbounded. `None` means "don't touch it", /// and candle's default is bounded, so a bad value degrades to bounded. #[test] - fn an_unparseable_value_does_not_become_unbounded() { + fn an_unparsable_value_does_not_become_unbounded() { for bad in ["", " ", "lots", "1GiB", "-1", "1.5"] { assert_eq!(parse_alloc_cache_capacity(Some(bad)), None, "{bad:?}"); } From d71fec0278b231770097655cd79b5b96333d7e5f Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Thu, 20 Aug 2026 18:39:27 +0100 Subject: [PATCH 18/27] fix(deps): regenerate Cargo.lock for the candle 89ab14ef pin `chore(deps): candle 89ab14ef` bumped all five candle crates in `Cargo.toml` but left `Cargo.lock` pinning 88d86a2. That is the stale-lock trap: a `--locked` CI build resolves the LOCK, so the lane would have compiled the OLD, UNBOUNDED allocator while reporting green on a PR whose entire subject is bounding it. `cargo metadata --locked` now succeeds; it failed before this commit. Regenerated with `cargo update -p candle-{core,nn,flash-attn,flash-attn-v3,metal-kernels}`. The only non-candle churn is a `windows-core` 0.61.2/0.62.2 dedup that fell out of the re-resolution; nothing arc builds on Linux or macOS reads it. --- Cargo.lock | 57 +++++++++++++++++++++--------------------------------- 1 file changed, 22 insertions(+), 35 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 1e5fa1e5f..5a68f4570 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -132,7 +132,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -143,7 +143,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -692,7 +692,7 @@ checksum = "ade8366b8bd5ba243f0a58f036cc0ca8a2f069cff1a2351ef1cac6b083e16fc0" [[package]] name = "candle-core" version = "0.9.2" -source = "git+https://github.com/aeonmindai/candle.git?rev=88d86a2019ee552923b217284e89847d2785dcdf#88d86a2019ee552923b217284e89847d2785dcdf" +source = "git+https://github.com/aeonmindai/candle.git?rev=89ab14ef1216331a782539f533b99f6816708604#89ab14ef1216331a782539f533b99f6816708604" dependencies = [ "accelerate-src", "byteorder", @@ -723,7 +723,7 @@ dependencies = [ [[package]] name = "candle-flash-attn" version = "0.9.2" -source = "git+https://github.com/aeonmindai/candle.git?rev=88d86a2019ee552923b217284e89847d2785dcdf#88d86a2019ee552923b217284e89847d2785dcdf" +source = "git+https://github.com/aeonmindai/candle.git?rev=89ab14ef1216331a782539f533b99f6816708604#89ab14ef1216331a782539f533b99f6816708604" dependencies = [ "anyhow", "candle-core", @@ -734,7 +734,7 @@ dependencies = [ [[package]] name = "candle-flash-attn-v3" version = "0.9.2" -source = "git+https://github.com/aeonmindai/candle.git?rev=88d86a2019ee552923b217284e89847d2785dcdf#88d86a2019ee552923b217284e89847d2785dcdf" +source = "git+https://github.com/aeonmindai/candle.git?rev=89ab14ef1216331a782539f533b99f6816708604#89ab14ef1216331a782539f533b99f6816708604" dependencies = [ "anyhow", "candle-core", @@ -747,7 +747,7 @@ dependencies = [ [[package]] name = "candle-kernels" version = "0.9.2" -source = "git+https://github.com/aeonmindai/candle.git?rev=88d86a2019ee552923b217284e89847d2785dcdf#88d86a2019ee552923b217284e89847d2785dcdf" +source = "git+https://github.com/aeonmindai/candle.git?rev=89ab14ef1216331a782539f533b99f6816708604#89ab14ef1216331a782539f533b99f6816708604" dependencies = [ "cudaforge", ] @@ -755,7 +755,7 @@ dependencies = [ [[package]] name = "candle-metal-kernels" version = "0.9.2" -source = "git+https://github.com/aeonmindai/candle.git?rev=88d86a2019ee552923b217284e89847d2785dcdf#88d86a2019ee552923b217284e89847d2785dcdf" +source = "git+https://github.com/aeonmindai/candle.git?rev=89ab14ef1216331a782539f533b99f6816708604#89ab14ef1216331a782539f533b99f6816708604" dependencies = [ "half", "objc2", @@ -769,7 +769,7 @@ dependencies = [ [[package]] name = "candle-nn" version = "0.9.2" -source = "git+https://github.com/aeonmindai/candle.git?rev=88d86a2019ee552923b217284e89847d2785dcdf#88d86a2019ee552923b217284e89847d2785dcdf" +source = "git+https://github.com/aeonmindai/candle.git?rev=89ab14ef1216331a782539f533b99f6816708604#89ab14ef1216331a782539f533b99f6816708604" dependencies = [ "accelerate-src", "candle-core", @@ -788,7 +788,7 @@ dependencies = [ [[package]] name = "candle-ug" version = "0.9.2" -source = "git+https://github.com/aeonmindai/candle.git?rev=88d86a2019ee552923b217284e89847d2785dcdf#88d86a2019ee552923b217284e89847d2785dcdf" +source = "git+https://github.com/aeonmindai/candle.git?rev=89ab14ef1216331a782539f533b99f6816708604#89ab14ef1216331a782539f533b99f6816708604" dependencies = [ "ug", "ug-cuda", @@ -1633,7 +1633,7 @@ dependencies = [ "libc", "option-ext", "redox_users 0.5.2", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -1813,7 +1813,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -2759,7 +2759,7 @@ dependencies = [ "js-sys", "log", "wasm-bindgen", - "windows-core 0.62.2", + "windows-core", ] [[package]] @@ -4091,7 +4091,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -5483,7 +5483,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.12.1", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -5542,7 +5542,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -6068,7 +6068,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -6515,7 +6515,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "230a1b821ccbd75b185820a1f1ff7b14d21da1e442e22c0863ea5f08771a8874" dependencies = [ "rustix 1.1.4", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -7626,7 +7626,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -7642,7 +7642,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9babd3a767a4c1aef6900409f85f5d53ce2544ccdfaa86dad48c91782c6d6893" dependencies = [ "windows-collections", - "windows-core 0.61.2", + "windows-core", "windows-future", "windows-link 0.1.3", "windows-numerics", @@ -7654,7 +7654,7 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3beeceb5e5cfd9eb1d76b381630e82c4241ccd0d27f1a39ed41b2760b255c5e8" dependencies = [ - "windows-core 0.61.2", + "windows-core", ] [[package]] @@ -7670,26 +7670,13 @@ dependencies = [ "windows-strings 0.4.2", ] -[[package]] -name = "windows-core" -version = "0.62.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" -dependencies = [ - "windows-implement", - "windows-interface", - "windows-link 0.2.1", - "windows-result 0.4.1", - "windows-strings 0.5.1", -] - [[package]] name = "windows-future" version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fc6a41e98427b19fe4b73c550f060b59fa592d7d686537eebf9385621bfbad8e" dependencies = [ - "windows-core 0.61.2", + "windows-core", "windows-link 0.1.3", "windows-threading", ] @@ -7734,7 +7721,7 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9150af68066c4c5c07ddc0ce30421554771e528bde427614c61038bc2c92c2b1" dependencies = [ - "windows-core 0.61.2", + "windows-core", "windows-link 0.1.3", ] From 9f110905b070e97977f91eef337b8a5fdf826730 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick Date: Wed, 19 Aug 2026 22:01:53 +0100 Subject: [PATCH 19/27] =?UTF-8?q?perf(ArcMoE):=20fuse=20the=20V4=20router?= =?UTF-8?q?=20region=20=E2=80=94=2018+9=20launches=20become=202,=20and=20f?= =?UTF-8?q?ix=20the=20sinkhorn=20local-memory=20spill?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit At b=1 the V4 decode step is launch-bound: ~7,924 cuLaunchKernel calls, of which the MoE router region is 3,259 (41.8%) across 86 contiguous spans — exactly two per layer (mhc_attn_pre, mhc_ffn_pre; 43 layers). Mean kernel in that region is 2.79 us and 89% are under 5 us, on [1,24] / [1,4,4] / [1,256] tensors. That is launch overhead wearing a kernel costume. Three changes, all bit-identical to what they replace: 1. cuda/hc_fused.cu `hc_pre_fused_f32` collapses hc_pre's 18-launch middle — a hand-decomposed RMS statistic in SEVEN launches (sqr, fast_sum, affine, affine, recip, sqrt, bmul) plus ELEVEN for the pre/post/comb scoring — into one kernel. `mixes` is never materialised. 2. cuda/hc_fused.cu `sqrt_softplus_f32` collapses MoeGate's sqrt(softplus(x)) from nine launches (zeros_like, bmaximum, uabs, uneg, uexp, affine, ulog, badd, usqrt) into one. 3. cuda/sinkhorn.cu is templated on `hc`. It took `hc` at runtime, so nvcc could not keep r[16]/buf[16]/col[16] in registers and demoted all three to local memory — `ptxas -v` reported 192 bytes stack frame. With hc=4 the kernel runs ONE block of FOUR threads, so nothing hides that latency: measured 30.0 us/call on a [1,4,4] tensor, 2.549 ms/step over 86 calls, 28% of the router's GPU time. Templating makes every index compile-time; ptxas now reports 0 bytes stack frame, 29 registers. Arithmetic unchanged. Bit-identity is by construction, not by tolerance: this region decides WHICH EXPERTS RUN, so a reassociated sum or a contracted FMA can change the emitted token. hc_fused.cu joins sinkhorn.cu in build.rs's dedicated no-fast-math, --fmad=false builder and carries the same #error guard, and it transcribes candle's ops exactly — including that `recipg` is `1.0 / a` with a *double* literal, that `mean_keepdim` is fast_sum followed by a separate affine, and that candle's FastReduce block_dim is min(1024, len).next_power_of_two(), which fixes the reduction tree's shape and therefore the f32 rounding. cuda/hc_fused.rs carries scalar replicas of both sides asserted bit-identical without a GPU, plus two anti-vacuity guards. The softplus one earned its place: its first version sampled [-30, -12.5, -1, 0.5, 3.25, 17] and passed against the naive log(1+exp(x)) form, because the two agree bit for bit until exp overflows near 88.7. It is now a three-way mutation probe. ARC_HC_FUSED=0 restores the eager chains so both can be A/B'd from one binary. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SpVNMpb13HkUXqSqbN1o9H --- mistralrs-core/build.rs | 27 +- mistralrs-core/src/cuda/ffi.rs | 28 ++ mistralrs-core/src/cuda/hc_fused.cu | 291 ++++++++++++ mistralrs-core/src/cuda/hc_fused.rs | 588 +++++++++++++++++++++++++ mistralrs-core/src/cuda/mod.rs | 1 + mistralrs-core/src/cuda/sinkhorn.cu | 203 ++++++--- mistralrs-core/src/models/deepseek4.rs | 20 +- mistralrs-core/src/models/dsv4_mhc.rs | 136 ++++-- 8 files changed, 1180 insertions(+), 114 deletions(-) create mode 100644 mistralrs-core/src/cuda/hc_fused.cu create mode 100644 mistralrs-core/src/cuda/hc_fused.rs diff --git a/mistralrs-core/build.rs b/mistralrs-core/build.rs index ef506a3db..bfe9f4226 100644 --- a/mistralrs-core/build.rs +++ b/mistralrs-core/build.rs @@ -10,15 +10,15 @@ 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. + // sinkhorn.cu and hc_fused.cu 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. Both sources carry an `#error` guard against + // being re-globbed under fast math. See mistralrs-core/src/cuda/. let mut builder = cudaforge::KernelBuilder::new() .source_glob("src/cuda/*.cu") - .exclude(&["sinkhorn.cu"]) + .exclude(&["sinkhorn.cu", "hc_fused.cu"]) .out_dir(&build_dir) .arg("-std=c++17") .arg("-O3") @@ -63,14 +63,15 @@ 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 - // mixes with the fast-math builder's. + // Dedicated IEEE (no fast math) builder for the bit-identity-critical + // kernels — bit-identity with candle-kernels requires accurate + // expf/logf + 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 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"]) + .source_files(vec!["src/cuda/sinkhorn.cu", "src/cuda/hc_fused.cu"]) .out_dir(&sinkhorn_dir) .arg("-std=c++17") .arg("-O3") diff --git a/mistralrs-core/src/cuda/ffi.rs b/mistralrs-core/src/cuda/ffi.rs index 81757e1b7..a2ca78e45 100644 --- a/mistralrs-core/src/cuda/ffi.rs +++ b/mistralrs-core/src/cuda/ffi.rs @@ -371,4 +371,32 @@ extern "C" { eps: f32, stream: i64, ); + // cuda/hc_fused.cu — fused V4 router-region kernels. Like sinkhorn.cu these + // are compiled by build.rs's dedicated no-fast-math builder; see the + // bit-identity contract at the top of that file. + #[allow(clippy::too_many_arguments)] + pub(crate) fn hc_pre_fused_f32( + x_flat: *const c_void, + mixes_raw: *const c_void, + hc_scale: *const c_void, + hc_base: *const c_void, + pre: *mut c_void, + post: *mut c_void, + comb_pre: *mut c_void, + n: i32, + d: i32, + m: i32, + hc: i32, + block_dim: i32, + inv_d: f32, + rms_eps: f32, + hc_eps: f32, + stream: i64, + ); + pub(crate) fn sqrt_softplus_f32( + inp: *const c_void, + out: *mut c_void, + numel: i64, + stream: i64, + ); } diff --git a/mistralrs-core/src/cuda/hc_fused.cu b/mistralrs-core/src/cuda/hc_fused.cu new file mode 100644 index 000000000..bdfc711ce --- /dev/null +++ b/mistralrs-core/src/cuda/hc_fused.cu @@ -0,0 +1,291 @@ +// Parent system: ArcInfer / ArcMoE +// +// Fused router-region kernels for DeepSeek-V4 (mHC pre-step + MoE gate scoring). +// +// --------------------------------------------------------------------------- +// WHY +// --------------------------------------------------------------------------- +// At b=1 the V4 decode step is launch-bound, not bandwidth-bound: the measured +// step issues ~7.9k `cuLaunchKernel` calls, of which the router region alone is +// 3,259 (41.8%) spread over 86 contiguous spans -- exactly two per layer +// (`mhc_attn_pre` and `mhc_ffn_pre`, 43 layers). Those spans operate on +// `[1, 24]` / `[1, 4, 4]` / `[1, 256]` tensors: mean kernel duration 2.79 us, +// 89% of them under 5 us. They are pure launch overhead wearing a kernel +// costume. +// +// Two expressions dominate the count and are collapsed here: +// +// 1. `hc_pre` (dsv4_mhc.rs) spells a hand-decomposed RMS statistic in SEVEN +// launches -- `sqr -> fast_sum -> affine(1/D) -> affine(+eps) -> urecip -> +// usqrt -> bmul` -- and then a further ELEVEN for the pre/post/comb +// sigmoid-and-bias scoring. Eighteen launches whose entire data footprint +// after the reduction is 24 floats. +// +// 2. `MoeGate::forward` (deepseek4.rs) spells `sqrt(softplus(x))` as NINE +// launches: `zeros_like -> bmaximum -> uabs -> uneg -> uexp -> affine(+1) +// -> ulog -> badd -> usqrt`. +// +// --------------------------------------------------------------------------- +// BIT-IDENTITY CONTRACT +// --------------------------------------------------------------------------- +// The mHC pre-step feeds `y` into attention and the gate scores decide WHICH +// EXPERTS RUN. A reassociated sum or a contracted FMA can flip an expert choice +// and therefore the generated token, so these kernels are bit-identical to the +// candle op chains they replace, BY CONSTRUCTION, not by tolerance. The same +// three rules that `sinkhorn.cu` documents apply verbatim: +// +// 1. REDUCTION ORDER. candle's `sum_keepdim` on CUDA runs candle-kernels +// `fast_sum` (reduce.cu) with `block_dim = min(1024, el_to_sum_per_block) +// .next_power_of_two()` (cuda_backend/mod.rs `FastReduce`). Each thread +// accumulates a *strided* slice sequentially into a zero-initialized +// accumulator (`shr[tid] = 0; shr[tid] += src[idx]; idx += blockDim.x`), +// then a pairwise tree `shr[t] += shr[t + s]` for s = block/2 .. 1. For +// D = hc_mult * hidden = 16384 that is 1024 threads x 16 elements each, and +// the tree order is NOT the sequential order. `hc_pre_fused_f32_kernel` +// replays it exactly, including the `shr[tid] = 0` initialisation (which +// canonicalises -0.0f to +0.0f) and the identity padding when a thread's +// strided slice is empty. +// +// 2. UNFUSED, IEEE ROUND-TO-NEAREST ARITHMETIC. candle-kernels build with +// plain -O3 and NO --use_fast_math, so their division is `div.rn.f32`, +// their `expf`/`logf` are the accurate libdevice `__nv_expf`/`__nv_logf`, +// and denormals are not flushed. mistralrs-core's build.rs compiles the +// rest of src/cuda/*.cu WITH --use_fast_math, which would silently rewrite +// all three. THIS FILE IS THEREFORE COMPILED BY THE SAME DEDICATED +// no-fast-math / --fmad=false builder that sinkhorn.cu uses, and carries +// the same #error guard against a future re-glob. All arithmetic uses +// __fadd_rn / __fmul_rn / __fsub_rn, which nvcc documents as never merged +// into an FMA, so the fusion cannot silently change rounding. +// +// 3. EXACT OP TRANSCRIPTION, including the ones that look like no-ops: +// - candle's `recipg(float a)` is `return 1.0 / a;` -- the literal is a +// *double*, so the op is `(float)(1.0 / (double)a)`, not `__frcp_rn(a)`. +// Transcribed literally in `candle_recip`. +// - `Tensor * f64` / `Tensor + f64` lower to `affine(mul, add)` whose +// kernel is `x * mul + add`, compiled by candle with the default +// -fmad=true, i.e. `fmaf(x, mul, add)`. Transcribed as explicit `fmaf` +// (an explicit fmaf is unaffected by our --fmad=false, which only +// governs *contraction* of separate mul+add). +// - `mean_keepdim` is `sum_impl(..)? * (1f64 / len as f64)`, i.e. a +// `fast_sum` followed by a SEPARATE affine -- the division is not folded +// into the reduction. +// - `candle_nn::ops::sigmoid` dispatches to candle-kernels `usigmoid_f32` +// = `sigmoid_fwd(x)` = `recipg(1.0f + expf(-x))`. +// +// Scalar Rust replicas of both sides (the candle op chain and this kernel) are +// asserted bit-identical over randomized inputs in `cuda/hc_fused.rs` +// `mod tests`, so a future edit to either side trips CPU CI without a GPU. The +// final proof is the on-GPU A/B: `ARC_HC_AB=1` recomputes the candle chain +// alongside the fused kernel and reports any bitwise mismatch. + +#include +#include + +#if defined(__USE_FAST_MATH__) +#error "hc_fused.cu must be compiled WITHOUT --use_fast_math: fast math rewrites expf/logf to the hardware approximations and IEEE division to approximate reciprocals, breaking bit-identity with candle-kernels (which build with plain -O3). See the dedicated no-fast-math builder in mistralrs-core/build.rs." +#endif + +namespace { + +// candle-kernels cuda_utils.cuh: +// __device__ __forceinline__ float recipg(float a) { return 1.0 / a; } +// The literal `1.0` is a double, so `a` is promoted, the division is done in +// double, and the result is narrowed on return. Transcribed literally rather +// than "simplified" to __frcp_rn: the two happen to agree for all finite +// inputs, but the point of this file is to not rely on such arguments. +__device__ __forceinline__ float candle_recip(float a) { + return (float)(1.0 / (double)a); +} + +// candle-kernels unary.cu: +// __device__ __forceinline__ T sigmoid_fwd(T x) { +// return recipg(static_cast(1) + expg(-x)); +// } +// with expg(float) == expf (libdevice __nv_expf under no-fast-math). +__device__ __forceinline__ float candle_sigmoid(float x) { + return candle_recip(__fadd_rn(1.0f, expf(-x))); +} + +} // namespace + +extern "C" { + +// --------------------------------------------------------------------------- +// hc_pre: the 18-launch RMS-scale + pre/post/comb scoring block, in one launch. +// --------------------------------------------------------------------------- +// +// Replaces, verbatim (dsv4_mhc.rs `hc_pre`): +// +// let sq_mean = x_flat.sqr()?.mean_keepdim(D::Minus1)?; // 3 +// let rsqrt = (sq_mean + rms_norm_eps)?.recip()?.sqrt()?; // 3 +// let mixes = mixes_raw.broadcast_mul(&rsqrt)?; // 1 +// let pre = sigmoid(pre_block *s_pre +b_pre)? + hc_eps; // 4 +// let post = sigmoid(post_block*s_post +b_post)?.affine(2,0); // 4 +// let comb_pre= comb_block*s_comb + b_comb; // 3 +// +// The trailing 3 for `comb` include the `.reshape((n, hc, hc))` of a narrowed +// (non-contiguous) view, which candle services with a real copy kernel; here +// the reshape is implicit in the output indexing and costs nothing. +// +// `mixes` itself is never materialised: it is consumed only by the three +// narrows, so the fused kernel keeps it in registers. +// +// Launch contract (enforced host-side in hc_fused.rs): +// grid = n blocks (one per row, matching candle's `dst_el` grid) +// block = min(1024, d).next_power_of_two() -- candle's FastReduce block_dim +// shmem = blockDim.x * sizeof(float) +// m <= blockDim.x, and x_flat / mixes_raw are contiguous F32. +__global__ void hc_pre_fused_f32_kernel( + const float *__restrict__ x_flat, // [n, d] F32 contiguous + const float *__restrict__ mixes_raw, // [n, m] F32 contiguous + const float *__restrict__ hc_scale, // [3] F32 + const float *__restrict__ hc_base, // [m] F32 + float *__restrict__ pre, // [n, hc] out + float *__restrict__ post, // [n, hc] out + float *__restrict__ comb_pre, // [n, hc, hc] out + int d, + int m, + int hc, + float inv_d, // (float)(1.0 / d) -- candle's mean scale, T::from_f64 + float rms_eps, // (float) rms_norm_eps + float hc_eps // (float) hc_eps +) { + extern __shared__ float shr[]; + const int n = blockIdx.x; + const int tid = threadIdx.x; + const int bd = blockDim.x; + + // ---- candle `sqr()` + `fast_sum` replay ------------------------------- + // candle runs these as two kernels; the intermediate `x*x` is rounded to + // f32 before the accumulation, which __fmul_rn reproduces. The strided + // walk and the zero-initialised accumulator are fast_sum's, verbatim. + const float *xrow = x_flat + (size_t)n * (size_t)d; + float acc = 0.0f; + for (int idx = tid; idx < d; idx += bd) { + const float v = xrow[idx]; + acc = __fadd_rn(acc, __fmul_rn(v, v)); + } + shr[tid] = acc; + // fast_sum's pairwise tree. The __syncthreads() sits at the TOP of the + // body there too, so the write above is visible before the first read. + for (int s = bd >> 1; s > 0; s >>= 1) { + __syncthreads(); + if (tid < s) { + shr[tid] = __fadd_rn(shr[tid], shr[tid + s]); + } + } + __syncthreads(); + + // ---- mean_keepdim tail, then `+eps -> recip -> sqrt` ------------------ + // Every thread recomputes this from shr[0]; it is deterministic scalar + // arithmetic on one value, so all threads agree bit for bit. + const float sum = shr[0]; + const float mean = fmaf(sum, inv_d, 0.0f); // mean_keepdim: affine(1/d, 0) + const float meps = fmaf(mean, 1.0f, rms_eps); // (sq_mean + rms_norm_eps) + const float rsqrt = sqrtf(candle_recip(meps)); + + // ---- broadcast_mul into `mixes`, then the three scoring blocks -------- + if (tid < m) { + const float mx = __fmul_rn(mixes_raw[(size_t)n * (size_t)m + tid], rsqrt); + if (tid < hc) { + // pre = sigmoid(pre_block * s_pre + b_pre) + hc_eps + const float t = __fadd_rn(__fmul_rn(mx, hc_scale[0]), hc_base[tid]); + pre[(size_t)n * (size_t)hc + tid] = fmaf(candle_sigmoid(t), 1.0f, hc_eps); + } else if (tid < 2 * hc) { + // post = 2 * sigmoid(post_block * s_post + b_post) + const int j = tid - hc; + const float t = __fadd_rn(__fmul_rn(mx, hc_scale[1]), hc_base[hc + j]); + post[(size_t)n * (size_t)hc + j] = fmaf(candle_sigmoid(t), 2.0f, 0.0f); + } else { + // comb_pre = comb_block * s_comb + b_comb (reshape is free here) + const int j = tid - 2 * hc; + const float t = __fadd_rn(__fmul_rn(mx, hc_scale[2]), hc_base[2 * hc + j]); + comb_pre[(size_t)n * (size_t)hc * (size_t)hc + j] = t; + } + } +} + +void hc_pre_fused_f32( + const void *x_flat, + const void *mixes_raw, + const void *hc_scale, + const void *hc_base, + void *pre, + void *post, + void *comb_pre, + int n, + int d, + int m, + int hc, + int block_dim, + float inv_d, + float rms_eps, + float hc_eps, + long long stream +) { + dim3 grid(n, 1, 1); + dim3 block(block_dim, 1, 1); + size_t shmem = (size_t)block_dim * sizeof(float); + hc_pre_fused_f32_kernel<<>>( + (const float *)x_flat, + (const float *)mixes_raw, + (const float *)hc_scale, + (const float *)hc_base, + (float *)pre, + (float *)post, + (float *)comb_pre, + d, + m, + hc, + inv_d, + rms_eps, + hc_eps); +} + +// --------------------------------------------------------------------------- +// sqrt(softplus(x)): the V4 gate scoring function, in one launch instead of 9. +// --------------------------------------------------------------------------- +// +// Replaces, verbatim (deepseek4.rs `MoeGate::forward`, ScoringFunc::SqrtSoftplus): +// +// let max0 = logits.maximum(&logits.zeros_like()?)?; // zeros + bmaximum +// let abs = logits.abs()?; // uabs +// let softplus = (max0 + ((abs.neg()?.exp()? + 1.0)?.log()?))?; +// softplus.sqrt()? +// +// candle op-for-op: bmaximum is `maxg(x, y)` == `fmaxf`; uabs is `fabsf`; +// uneg is `-x`; uexp is `expf`; `+ 1.0` is `affine(1.0, 1.0)` == `fmaf(x,1,1)`; +// ulog is `logf`; the outer `+` is `badd` == `x + y`; usqrt is `sqrtf`. +__global__ void sqrt_softplus_f32_kernel( + const float *__restrict__ inp, + float *__restrict__ out, + long long numel +) { + const long long stride = (long long)blockDim.x * (long long)gridDim.x; + for (long long i = (long long)blockIdx.x * (long long)blockDim.x + threadIdx.x; + i < numel; + i += stride) { + const float l = inp[i]; + const float mx = fmaxf(l, 0.0f); // maximum(logits, zeros_like) + const float a = fabsf(l); // abs() + const float e = expf(-a); // neg() then exp() + const float p1 = fmaf(e, 1.0f, 1.0f); // + 1.0 (affine) + const float lg = logf(p1); // log() + out[i] = sqrtf(__fadd_rn(mx, lg)); // max0 + ... then sqrt() + } +} + +void sqrt_softplus_f32(const void *inp, void *out, long long numel, long long stream) { + const int block = 256; + long long blocks = (numel + block - 1) / block; + if (blocks < 1) { + blocks = 1; + } + if (blocks > 65535) { + blocks = 65535; + } + sqrt_softplus_f32_kernel<<<(unsigned)blocks, block, 0, (cudaStream_t)stream>>>( + (const float *)inp, (float *)out, numel); +} + +} // extern "C" diff --git a/mistralrs-core/src/cuda/hc_fused.rs b/mistralrs-core/src/cuda/hc_fused.rs new file mode 100644 index 000000000..c176e93ea --- /dev/null +++ b/mistralrs-core/src/cuda/hc_fused.rs @@ -0,0 +1,588 @@ +//! Parent system: ArcInfer / ArcMoE +//! +//! Rust side of `cuda/hc_fused.cu` — the fused DeepSeek-V4 router-region +//! kernels. +//! +//! Two entry points, each replacing a launch-bound candle op chain with one +//! kernel: +//! +//! - [`hc_pre_fused_cuda`] — the 18-launch RMS-scale + pre/post/comb scoring +//! block of [`crate::models::dsv4_mhc::V4MHCLayerParams::hc_pre`]. +//! - [`sqrt_softplus_cuda`] — the 9-launch `sqrt(softplus(x))` gate scoring +//! function of `MoeGate::forward`. +//! +//! Both are **bit-identical** to the chains they replace, not approximations. +//! The contract, and the three ways candle's arithmetic can be got wrong, are +//! documented at the top of `hc_fused.cu`. [`reference`] below carries scalar +//! replicas of both sides that `mod tests` asserts bit-identical without a GPU. +//! +//! Both are opt-out via `ARC_HC_FUSED=0`, which restores the eager chain. That +//! switch exists so the on-GPU A/B (fused vs. eager, greedy decode, compare +//! tokens) can be run from one binary. + +/// `ARC_HC_FUSED=0` disables the fused kernels and restores the eager candle +/// chains. Any other value (or unset) keeps them on. +/// +/// Read once and cached: this is consulted on every layer of every decode step, +/// and `std::env::var` takes a global lock. +pub fn fused_enabled() -> bool { + use std::sync::OnceLock; + static ENABLED: OnceLock = OnceLock::new(); + *ENABLED.get_or_init(|| !matches!(std::env::var("ARC_HC_FUSED").as_deref(), Ok("0"))) +} + +/// candle's `FastReduce` block size for a reduction of `len` elements +/// (`cuda_backend/mod.rs`: `usize::min(1024, el_to_sum_per_block) +/// .next_power_of_two()`). The fused kernel must use exactly this, because the +/// reduction tree's shape — and therefore the f32 rounding — depends on it. +pub(crate) fn candle_reduce_block_dim(len: usize) -> usize { + usize::min(1024, len).next_power_of_two() +} + +#[cfg(feature = "cuda")] +mod cuda_impl { + use candle_core as candle; + use candle_core::{DType, Result, Tensor}; + + use super::candle_reduce_block_dim; + + /// Pull a contiguous F32 CUDA tensor's device pointer. + fn f32_ptr(t: &Tensor, what: &str) -> Result<*const std::ffi::c_void> { + use candle_core::cuda_backend::cudarc::driver::DevicePtr; + if t.dtype() != DType::F32 { + candle::bail!("hc_fused: {what} must be F32, got {:?}", t.dtype()); + } + if !t.is_contiguous() { + candle::bail!("hc_fused: {what} must be contiguous"); + } + let (s, l) = t.storage_and_layout(); + let s = match &*s { + candle::Storage::Cuda(c) => c.as_cuda_slice::()?, + _ => candle::bail!("hc_fused: {what} must be on CUDA"), + }; + Ok(s.slice(l.start_offset()..).device_ptr(s.stream()).0 as *const std::ffi::c_void) + } + + /// Fused `hc_pre` middle section: the RMS statistic, its broadcast into + /// `mixes`, and the three scoring blocks. + /// + /// Inputs (all F32, contiguous, same CUDA device): + /// - `x_flat` `[n, d]` — the promoted residual stack, `d = hc * hidden` + /// - `mixes_raw` `[n, m]` — the gate GEMM output, `m = (2 + hc) * hc` + /// - `hc_scale` `[3]` + /// - `hc_base` `[m]` + /// + /// Returns `(pre [n, hc], post [n, hc], comb_pre [n, hc, hc])`. + #[allow(clippy::too_many_arguments)] + pub fn hc_pre_fused_cuda( + x_flat: &Tensor, + mixes_raw: &Tensor, + hc_scale: &Tensor, + hc_base: &Tensor, + hc: usize, + rms_eps: f64, + hc_eps: f64, + ) -> Result<(Tensor, Tensor, Tensor)> { + use candle_core::cuda_backend::cudarc::driver::DevicePtr; + + let (n, d) = x_flat.dims2()?; + let (n2, m) = mixes_raw.dims2()?; + if n != n2 { + candle::bail!("hc_fused: x_flat has {n} rows but mixes_raw has {n2}"); + } + if m != (2 + hc) * hc { + candle::bail!( + "hc_fused: mixes_raw has {m} columns, expected (2 + hc) * hc = {} for hc={hc}", + (2 + hc) * hc + ); + } + if hc_scale.dims1()? != 3 { + candle::bail!("hc_fused: hc_scale must be [3]"); + } + if hc_base.dims1()? != m { + candle::bail!("hc_fused: hc_base must be [{m}]"); + } + let block_dim = candle_reduce_block_dim(d); + if m > block_dim { + // The scoring tail is done by threads 0..m, so it must fit the + // block candle's reduction shape dictates. Unreachable for V4 + // (m = 24, block_dim = 1024) but a wrong answer if it ever isn't. + candle::bail!( + "hc_fused: m={m} exceeds the reduction block_dim={block_dim} implied by d={d}" + ); + } + if n == 0 { + candle::bail!("hc_fused: empty batch"); + } + + let x_ptr = f32_ptr(x_flat, "x_flat")?; + let mixes_ptr = f32_ptr(mixes_raw, "mixes_raw")?; + let scale_ptr = f32_ptr(hc_scale, "hc_scale")?; + let base_ptr = f32_ptr(hc_base, "hc_base")?; + + let dev = x_flat.device().as_cuda_device()?; + let pre_buf = unsafe { dev.alloc::(n * hc) }?; + let post_buf = unsafe { dev.alloc::(n * hc) }?; + let comb_buf = unsafe { dev.alloc::(n * hc * hc) }?; + let stream = dev.cuda_stream().cu_stream() as i64; + + #[allow(clippy::cast_possible_truncation)] + unsafe { + crate::cuda::ffi::hc_pre_fused_f32( + x_ptr, + mixes_ptr, + scale_ptr, + base_ptr, + pre_buf.device_ptr(pre_buf.stream()).0 as *mut std::ffi::c_void, + post_buf.device_ptr(post_buf.stream()).0 as *mut std::ffi::c_void, + comb_buf.device_ptr(comb_buf.stream()).0 as *mut std::ffi::c_void, + n as i32, + d as i32, + m as i32, + hc as i32, + block_dim as i32, + // candle's `mean_keepdim` scale is `T::from_f64(1f64 / len)`, + // i.e. the f64 reciprocal narrowed to f32 — not `1.0f32 / d`. + (1f64 / d as f64) as f32, + rms_eps as f32, + hc_eps as f32, + stream, + ); + } + + let wrap = |buf, shape| { + let st = candle::CudaStorage::wrap_cuda_slice(buf, dev.clone()); + Tensor::from((candle::Storage::Cuda(st), shape)) + }; + Ok(( + wrap(pre_buf, (n, hc)), + wrap(post_buf, (n, hc)), + wrap(comb_buf, (n, hc, hc)), + )) + } + + /// Fused `sqrt(softplus(x))` — the V4 gate scoring function. + pub fn sqrt_softplus_cuda(logits: &Tensor) -> Result { + use candle_core::cuda_backend::cudarc::driver::DevicePtr; + + let logits = logits.contiguous()?; + let numel = logits.elem_count(); + if numel == 0 { + candle::bail!("hc_fused: sqrt_softplus on an empty tensor"); + } + let in_ptr = f32_ptr(&logits, "logits")?; + let dev = logits.device().as_cuda_device()?; + let out_buf = unsafe { dev.alloc::(numel) }?; + let stream = dev.cuda_stream().cu_stream() as i64; + + #[allow(clippy::cast_possible_truncation)] + unsafe { + crate::cuda::ffi::sqrt_softplus_f32( + in_ptr, + out_buf.device_ptr(out_buf.stream()).0 as *mut std::ffi::c_void, + numel as i64, + stream, + ); + } + + let st = candle::CudaStorage::wrap_cuda_slice(out_buf, dev.clone()); + Ok(Tensor::from(( + candle::Storage::Cuda(st), + logits.shape().clone(), + ))) + } +} + +#[cfg(feature = "cuda")] +pub use cuda_impl::{hc_pre_fused_cuda, sqrt_softplus_cuda}; + +// Non-CUDA stubs so call sites need no `cfg`. They are never reached: every +// caller gates on [`usable`], which is false without a CUDA device. +#[cfg(not(feature = "cuda"))] +#[allow(clippy::too_many_arguments)] +pub fn hc_pre_fused_cuda( + _x_flat: &candle_core::Tensor, + _mixes_raw: &candle_core::Tensor, + _hc_scale: &candle_core::Tensor, + _hc_base: &candle_core::Tensor, + _hc: usize, + _rms_eps: f64, + _hc_eps: f64, +) -> candle_core::Result<( + candle_core::Tensor, + candle_core::Tensor, + candle_core::Tensor, +)> { + candle_core::bail!("hc_pre_fused_cuda requires the cuda feature") +} + +#[cfg(not(feature = "cuda"))] +pub fn sqrt_softplus_cuda(_logits: &candle_core::Tensor) -> candle_core::Result { + candle_core::bail!("sqrt_softplus_cuda requires the cuda feature") +} + +/// Whether the fused path should be taken for a tensor: a CUDA device, F32, +/// and not disabled by `ARC_HC_FUSED=0`. +/// +/// Checked BEFORE the call rather than by catching an error from it, so that a +/// genuine CUDA failure inside the kernel surfaces as an error instead of being +/// silently swallowed into the eager path (the "silent success" failure mode +/// this codebase has been bitten by before). +pub fn usable(t: &candle_core::Tensor) -> bool { + cfg!(feature = "cuda") + && t.device().is_cuda() + && t.dtype() == candle_core::DType::F32 + && fused_enabled() +} + +/// Scalar f32 replicas that pin the fused kernels' op ORDER and ROUNDING to +/// what the candle CUDA backend actually executes, bit for bit. +/// +/// Same discipline, and same known scalar-vs-GPU gap, as +/// [`crate::cuda::sinkhorn::reference`]: Rust/libm `f32::exp`/`ln` may differ +/// from CUDA libdevice `__nv_expf`/`__nv_logf` in the last ulp, but that +/// difference CANCELS in the on-GPU A/B because the candle chain and the fused +/// kernel call the same libdevice routine. Everything else here — add, mul, +/// div, fma, sqrt, single rounding per op — is exact IEEE f32 on both sides, so +/// the bitwise assertions in `mod tests` are meaningful. +#[allow(dead_code)] // consumed by `mod tests` +pub(crate) mod reference { + /// candle-kernels `cuda_utils.cuh`: `recipg(float a) { return 1.0 / a; }`. + /// The literal is a double, so this is `(float)(1.0 / (double)a)`. + pub(crate) fn candle_recip(a: f32) -> f32 { + (1.0f64 / a as f64) as f32 + } + + /// candle-kernels `unary.cu`: `sigmoid_fwd(x) = recipg(1 + expg(-x))`. + pub(crate) fn candle_sigmoid(x: f32) -> f32 { + candle_recip(1.0f32 + (-x).exp()) + } + + /// candle-kernels `fast_sum` (reduce.cu) over a contiguous row, with + /// candle's `FastReduce` block size. Each virtual thread accumulates its + /// strided slice sequentially into a zero-initialised accumulator, then the + /// pairwise tree runs over `block_dim` slots. + pub(crate) fn candle_fast_sum(vals: &[f32], block_dim: usize) -> f32 { + let mut shr = vec![0.0f32; block_dim]; + for (tid, slot) in shr.iter_mut().enumerate() { + let mut idx = tid; + while idx < vals.len() { + *slot += vals[idx]; + idx += block_dim; + } + } + let mut s = block_dim / 2; + while s > 0 { + for t in 0..s { + shr[t] += shr[t + s]; + } + s /= 2; + } + shr[0] + } + + /// The **candle op chain** of `dsv4_mhc::hc_pre`, one step per kernel + /// launch, in the order the eager path issues them. + /// + /// Returns `(pre, post, comb_pre)`. + #[allow(clippy::type_complexity)] + pub(crate) fn hc_pre_candle_replay( + x_flat: &[f32], + mixes_raw: &[f32], + hc_scale: &[f32], + hc_base: &[f32], + hc: usize, + rms_eps: f64, + hc_eps: f64, + ) -> (Vec, Vec, Vec) { + let d = x_flat.len(); + let m = mixes_raw.len(); + let block_dim = super::candle_reduce_block_dim(d); + + // sqr() -> a whole separate kernel, so x*x is rounded to f32 first. + let sq: Vec = x_flat.iter().map(|v| v * v).collect(); + // mean_keepdim = fast_sum, then a SEPARATE affine(1/d, 0) == fmaf. + let sum = candle_fast_sum(&sq, block_dim); + let inv_d = (1f64 / d as f64) as f32; + let mean = sum.mul_add(inv_d, 0.0f32); + // + rms_norm_eps -> affine(1.0, eps) == fmaf(x, 1, eps) + let meps = mean.mul_add(1.0f32, rms_eps as f32); + // recip() then sqrt() + let rsqrt = candle_recip(meps).sqrt(); + + // broadcast_mul -> `mixes` + let mixes: Vec = mixes_raw.iter().map(|v| v * rsqrt).collect(); + + let mut pre = Vec::with_capacity(hc); + for j in 0..hc { + // broadcast_mul, broadcast_add, sigmoid, affine(1, hc_eps) + let t = mixes[j] * hc_scale[0] + hc_base[j]; + pre.push(candle_sigmoid(t).mul_add(1.0f32, hc_eps as f32)); + } + let mut post = Vec::with_capacity(hc); + for j in 0..hc { + let t = mixes[hc + j] * hc_scale[1] + hc_base[hc + j]; + post.push(candle_sigmoid(t).mul_add(2.0f32, 0.0f32)); + } + let mut comb = Vec::with_capacity(hc * hc); + for j in 0..(m - 2 * hc) { + comb.push(mixes[2 * hc + j] * hc_scale[2] + hc_base[2 * hc + j]); + } + (pre, post, comb) + } + + /// The **fused kernel** of `hc_fused.cu`, transcribed: one block, `acc` in + /// a register, the same tree, then the scoring tail. + #[allow(clippy::type_complexity)] + pub(crate) fn hc_pre_fused_replay( + x_flat: &[f32], + mixes_raw: &[f32], + hc_scale: &[f32], + hc_base: &[f32], + hc: usize, + rms_eps: f64, + hc_eps: f64, + ) -> (Vec, Vec, Vec) { + let d = x_flat.len(); + let m = mixes_raw.len(); + let block_dim = super::candle_reduce_block_dim(d); + + // Per-thread: acc = fadd(acc, fmul(v, v)) over the strided slice. + let mut shr = vec![0.0f32; block_dim]; + for (tid, slot) in shr.iter_mut().enumerate() { + let mut acc = 0.0f32; + let mut idx = tid; + while idx < d { + let v = x_flat[idx]; + acc += v * v; + idx += block_dim; + } + *slot = acc; + } + let mut s = block_dim / 2; + while s > 0 { + for t in 0..s { + shr[t] += shr[t + s]; + } + s /= 2; + } + let sum = shr[0]; + let inv_d = (1f64 / d as f64) as f32; + let mean = sum.mul_add(inv_d, 0.0f32); + let meps = mean.mul_add(1.0f32, rms_eps as f32); + let rsqrt = candle_recip(meps).sqrt(); + + let mut pre = vec![0.0f32; hc]; + let mut post = vec![0.0f32; hc]; + let mut comb = vec![0.0f32; m - 2 * hc]; + for tid in 0..m { + let mx = mixes_raw[tid] * rsqrt; + if tid < hc { + let t = mx * hc_scale[0] + hc_base[tid]; + pre[tid] = candle_sigmoid(t).mul_add(1.0f32, hc_eps as f32); + } else if tid < 2 * hc { + let j = tid - hc; + let t = mx * hc_scale[1] + hc_base[hc + j]; + post[j] = candle_sigmoid(t).mul_add(2.0f32, 0.0f32); + } else { + let j = tid - 2 * hc; + let t = mx * hc_scale[2] + hc_base[2 * hc + j]; + comb[j] = t; + } + } + (pre, post, comb) + } + + /// The **candle op chain** of `MoeGate::forward`'s `SqrtSoftplus` arm. + pub(crate) fn sqrt_softplus_candle_replay(logits: &[f32]) -> Vec { + // max0 = maximum(logits, zeros_like) + let max0: Vec = logits.iter().map(|&l| l.max(0.0f32)).collect(); + // abs -> neg -> exp -> affine(1, 1) -> log + let a: Vec = logits.iter().map(|&l| l.abs()).collect(); + let n: Vec = a.iter().map(|&v| -v).collect(); + let e: Vec = n.iter().map(|&v| v.exp()).collect(); + let p1: Vec = e.iter().map(|&v| v.mul_add(1.0f32, 1.0f32)).collect(); + let lg: Vec = p1.iter().map(|&v| v.ln()).collect(); + // badd then usqrt + max0.iter() + .zip(lg.iter()) + .map(|(&x, &y)| (x + y).sqrt()) + .collect() + } + + /// The **fused kernel** of `hc_fused.cu`, transcribed. + pub(crate) fn sqrt_softplus_fused_replay(logits: &[f32]) -> Vec { + logits + .iter() + .map(|&l| { + let mx = l.max(0.0f32); + let a = l.abs(); + let e = (-a).exp(); + let p1 = e.mul_add(1.0f32, 1.0f32); + let lg = p1.ln(); + (mx + lg).sqrt() + }) + .collect() + } +} + +#[cfg(test)] +mod tests { + use super::reference::*; + + /// Deterministic LCG so a failure is reproducible without a fixture file. + struct Lcg(u64); + impl Lcg { + fn next_f32(&mut self, lo: f32, hi: f32) -> f32 { + self.0 = self.0.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407); + let u = ((self.0 >> 40) as f32) / ((1u32 << 24) as f32); + lo + u * (hi - lo) + } + } + + #[test] + fn candle_reduce_block_dim_matches_fast_reduce() { + // cuda_backend/mod.rs: usize::min(1024, el_to_sum_per_block).next_power_of_two() + assert_eq!(super::candle_reduce_block_dim(16384), 1024); // V4: hc_mult * hidden + assert_eq!(super::candle_reduce_block_dim(1024), 1024); + assert_eq!(super::candle_reduce_block_dim(1000), 1024); + assert_eq!(super::candle_reduce_block_dim(24), 32); + assert_eq!(super::candle_reduce_block_dim(4), 4); + assert_eq!(super::candle_reduce_block_dim(3), 4); + } + + /// The whole point of the file: the fused kernel and the candle op chain + /// must agree BITWISE, at V4's real shapes. + #[test] + fn hc_pre_fused_is_bit_identical_to_candle_chain() { + let hc = 4usize; + let hidden = 4096usize; + let d = hc * hidden; + let m = (2 + hc) * hc; + let mut rng = Lcg(0x5eed_1234_abcd_0001); + + for trial in 0..8 { + let x: Vec = (0..d).map(|_| rng.next_f32(-3.0, 3.0)).collect(); + let mixes: Vec = (0..m).map(|_| rng.next_f32(-6.0, 6.0)).collect(); + let scale: Vec = (0..3).map(|_| rng.next_f32(-2.0, 2.0)).collect(); + let base: Vec = (0..m).map(|_| rng.next_f32(-2.0, 2.0)).collect(); + + let (pa, qa, ca) = hc_pre_candle_replay(&x, &mixes, &scale, &base, hc, 1e-6, 1e-6); + let (pb, qb, cb) = hc_pre_fused_replay(&x, &mixes, &scale, &base, hc, 1e-6, 1e-6); + + for (i, (a, b)) in pa.iter().zip(pb.iter()).enumerate() { + assert_eq!(a.to_bits(), b.to_bits(), "trial {trial} pre[{i}]: {a} vs {b}"); + } + for (i, (a, b)) in qa.iter().zip(qb.iter()).enumerate() { + assert_eq!(a.to_bits(), b.to_bits(), "trial {trial} post[{i}]: {a} vs {b}"); + } + for (i, (a, b)) in ca.iter().zip(cb.iter()).enumerate() { + assert_eq!(a.to_bits(), b.to_bits(), "trial {trial} comb[{i}]: {a} vs {b}"); + } + } + } + + /// A guard that cannot go red is not a guard. Reassociating the reduction — + /// the single most likely way to break this fusion — must be detected. + #[test] + fn hc_pre_guard_detects_a_reassociated_reduction() { + let hc = 4usize; + let d = hc * 4096; + let m = (2 + hc) * hc; + let mut rng = Lcg(0xdead_beef_0000_0007); + let x: Vec = (0..d).map(|_| rng.next_f32(-3.0, 3.0)).collect(); + let mixes: Vec = (0..m).map(|_| rng.next_f32(-6.0, 6.0)).collect(); + let scale: Vec = (0..3).map(|_| rng.next_f32(-2.0, 2.0)).collect(); + let base: Vec = (0..m).map(|_| rng.next_f32(-2.0, 2.0)).collect(); + + let (pre_ref, ..) = hc_pre_candle_replay(&x, &mixes, &scale, &base, hc, 1e-6, 1e-6); + + // Naive left-to-right sum instead of candle's strided-then-tree order. + let mut seq = 0.0f32; + for v in &x { + seq += v * v; + } + let inv_d = (1f64 / d as f64) as f32; + let rsqrt = candle_recip(seq.mul_add(inv_d, 0.0f32).mul_add(1.0f32, 1e-6f32)).sqrt(); + let pre_wrong: Vec = (0..hc) + .map(|j| { + let t = (mixes[j] * rsqrt) * scale[0] + base[j]; + candle_sigmoid(t).mul_add(1.0f32, 1e-6f32) + }) + .collect(); + + assert!( + pre_ref + .iter() + .zip(pre_wrong.iter()) + .any(|(a, b)| a.to_bits() != b.to_bits()), + "the reduction-order guard is vacuous: a sequential sum of {d} squares produced \ + bit-identical output to candle's strided+tree order, so this test would pass on a \ + kernel that reassociates the reduction" + ); + } + + #[test] + fn sqrt_softplus_fused_is_bit_identical_to_candle_chain() { + let mut rng = Lcg(0x0bad_c0de_1111_2222); + // n_routed_experts = 256 for V4-Flash. + let logits: Vec = (0..256).map(|_| rng.next_f32(-40.0, 40.0)).collect(); + let a = sqrt_softplus_candle_replay(&logits); + let b = sqrt_softplus_fused_replay(&logits); + for (i, (x, y)) in a.iter().zip(b.iter()).enumerate() { + assert_eq!(x.to_bits(), y.to_bits(), "logit[{i}]={}: {x} vs {y}", logits[i]); + } + // Including the exact-zero / sign boundary that `maximum(x, 0)` and + // `abs(x)` disagree on. + let edge = [0.0f32, -0.0f32, 1e-30, -1e-30, 88.0, -88.0, 1.0, -1.0]; + let a = sqrt_softplus_candle_replay(&edge); + let b = sqrt_softplus_fused_replay(&edge); + for (i, (x, y)) in a.iter().zip(b.iter()).enumerate() { + assert_eq!(x.to_bits(), y.to_bits(), "edge[{i}]={}: {x} vs {y}", edge[i]); + } + } + + /// Same discipline for the softplus guard: prove it can go red. + /// + /// Mutation-tests the three ways a careless reimplementation of + /// `sqrt(softplus(x))` actually goes wrong, and asserts the bitwise + /// comparison catches each. The first one is the reason this test exists at + /// all: an earlier version of it sampled only `[-30, -12.5, -1, 0.5, 3.25, + /// 17]` and PASSED against the naive form, because over that range the two + /// formulations happen to agree bit for bit — the stable split only earns + /// its keep once `exp(x)` overflows f32 near x = 88.7. A guard that green- + /// lights the unstable form is not a guard. + #[test] + fn sqrt_softplus_guard_detects_plausible_mistranscriptions() { + // Spans the overflow boundary (exp overflows f32 above ~88.7) and the + // sign boundary that `max(x, 0)` and `-|x|` turn on. + let logits: Vec = vec![-120.0, -30.0, -1.0, -0.25, 0.0, 0.25, 3.25, 17.0, 95.0, 200.0]; + let reference = sqrt_softplus_candle_replay(&logits); + + let mutants: [(&str, fn(f32) -> f32); 3] = [ + // 1. the unstable form: log(1 + exp(x)), no max/abs split. + ("unstable log(1+exp(x))", |l| (1.0f32 + l.exp()).ln().sqrt()), + // 2. dropped the max(x, 0) term. + ("dropped max(x,0)", |l| { + (1.0f32 + (-l.abs()).exp()).ln().sqrt() + }), + // 3. lost the negation on the exponent. + ("exp(+|x|) instead of exp(-|x|)", |l| { + (l.max(0.0f32) + (1.0f32 + l.abs().exp()).ln()).sqrt() + }), + ]; + + for (name, mutate) in mutants { + let got: Vec = logits.iter().map(|&l| mutate(l)).collect(); + assert!( + reference + .iter() + .zip(got.iter()) + .any(|(a, b)| a.to_bits() != b.to_bits()), + "the softplus guard is vacuous for mutation '{name}': it produced bit-identical \ + output to the reference on every sample, so this test would pass on a kernel \ + carrying that bug" + ); + } + } +} diff --git a/mistralrs-core/src/cuda/mod.rs b/mistralrs-core/src/cuda/mod.rs index 7318477f8..68f0e65ff 100644 --- a/mistralrs-core/src/cuda/mod.rs +++ b/mistralrs-core/src/cuda/mod.rs @@ -1,5 +1,6 @@ pub mod ffi; pub mod gdn; +pub mod hc_fused; pub mod moe; pub mod sinkhorn; pub mod ssm; diff --git a/mistralrs-core/src/cuda/sinkhorn.cu b/mistralrs-core/src/cuda/sinkhorn.cu index ae665e64b..767ac55a1 100644 --- a/mistralrs-core/src/cuda/sinkhorn.cu +++ b/mistralrs-core/src/cuda/sinkhorn.cu @@ -70,123 +70,195 @@ // hc_mult is 4 for V4-Flash; cap at 16 for safety (shared-mem sized at launch). #define SINKHORN_MAX_HC 16 +// --------------------------------------------------------------------------- +// PERFORMANCE NOTE (why this file is templated on `hc`) +// --------------------------------------------------------------------------- +// The first version of this kernel took `hc` as a RUNTIME argument and held the +// per-thread row / tree buffers in `float r[SINKHORN_MAX_HC]` arrays walked by +// runtime-bounded loops. nvcc cannot keep an array in registers unless every +// index is resolvable at compile time, so all three arrays were demoted to +// LOCAL MEMORY -- `ptxas -v` reported `192 bytes stack frame` (48 floats = +// r[16] + buf[16] + col[16]). With hc = 4 the kernel launches ONE block of FOUR +// threads, so there is no occupancy to hide that latency behind: every one of +// the ~1,800 dependent local-memory round trips per call (20 iterations x two +// tree reductions x two arrays, plus the row read/write) is exposed. Measured +// cost: 30.0 us per call on a [1, 4, 4] tensor -- 2.549 ms/step over the 86 +// calls, 28% of the whole router region's GPU time, for 16 floats of work. +// +// Templating on `hc` makes every index a compile-time constant, the arrays +// become registers (`0 bytes stack frame`), and the arithmetic is UNCHANGED -- +// same ops, same order, same intrinsics -- so bit-identity is preserved by +// construction. The tree reductions are expressed as template recursion rather +// than `#pragma unroll` loops specifically so that the halving step `S` is a +// constant in the type system and cannot silently fall back to a runtime index. +// --------------------------------------------------------------------------- + namespace { -__device__ __forceinline__ int next_pow2_le16(int v) { +constexpr int next_pow2_ce(int v) { int p = 1; while (p < v) p <<= 1; return p; } +// One level of candle's pairwise tree, with the stride `S` fixed by the type. +// `buf` is taken by array reference (not pointer) so the indices stay +// compile-time and the storage stays in registers. +template struct TreeSumLevel { + __device__ __forceinline__ static void run(float (&buf)[N]) { + constexpr int S = P / 2; +#pragma unroll + for (int t = 0; t < S; ++t) { + buf[t] = __fadd_rn(buf[t], buf[t + S]); + } + TreeSumLevel::run(buf); + } +}; +template struct TreeSumLevel<1, N> { + __device__ __forceinline__ static void run(float (&)[N]) {} +}; + +template struct TreeMaxLevel { + __device__ __forceinline__ static void run(float (&buf)[N]) { + constexpr int S = P / 2; +#pragma unroll + for (int t = 0; t < S; ++t) { + buf[t] = fmaxf(buf[t], buf[t + S]); + } + TreeMaxLevel::run(buf); + } +}; +template struct TreeMaxLevel<1, N> { + __device__ __forceinline__ static void run(float (&)[N]) {} +}; + // Replays candle-kernels `fast_sum` (reduce.cu) for reduced_len <= 16: // zero-initialized accumulators, one element per virtual thread // (block_dim = next_pow2(len) >= len so each thread loads at most one), // then the pairwise tree. `__fadd_rn(0.0f, v)` mirrors `shr[tid] = 0; // shr[tid] += v` (note: turns -0.0f into +0.0f, exactly like candle). -__device__ __forceinline__ float candle_tree_sum(const float* v, int len) { - float buf[SINKHORN_MAX_HC]; - const int p = next_pow2_le16(len); - for (int t = 0; t < p; ++t) { - buf[t] = (t < len) ? __fadd_rn(0.0f, v[t]) : 0.0f; - } - for (int s = p >> 1; s > 0; s >>= 1) { - for (int t = 0; t < s; ++t) { - buf[t] = __fadd_rn(buf[t], buf[t + s]); - } +template __device__ __forceinline__ float candle_tree_sum(const float (&v)[HC]) { + constexpr int P = next_pow2_ce(HC); + float buf[P]; +#pragma unroll + for (int t = 0; t < P; ++t) { + buf[t] = (t < HC) ? __fadd_rn(0.0f, v[t]) : 0.0f; } + TreeSumLevel::run(buf); return buf[0]; } // Replays candle-kernels `fast_max` (reduce.cu): -INF init, maxg == fmaxf // (NaN-ignoring IEEE maxNum), pairwise tree. Order-insensitive for finite // inputs but mirrored anyway so NaN propagation matches candle exactly. -__device__ __forceinline__ float candle_tree_max(const float* v, int len) { - float buf[SINKHORN_MAX_HC]; - const int p = next_pow2_le16(len); - for (int t = 0; t < p; ++t) { - buf[t] = (t < len) ? fmaxf(-INFINITY, v[t]) : -INFINITY; - } - for (int s = p >> 1; s > 0; s >>= 1) { - for (int t = 0; t < s; ++t) { - buf[t] = fmaxf(buf[t], buf[t + s]); - } +template __device__ __forceinline__ float candle_tree_max(const float (&v)[HC]) { + constexpr int P = next_pow2_ce(HC); + float buf[P]; +#pragma unroll + for (int t = 0; t < P; ++t) { + buf[t] = (t < HC) ? fmaxf(-INFINITY, v[t]) : -INFINITY; } + TreeMaxLevel::run(buf); return buf[0]; } -} // namespace - -extern "C" { - // in/out: [n, hc, hc] row-major F32. One block per matrix `n`, `hc` threads. -// Shared memory: hc*hc (matrix tile) + hc (column sums) floats. +// `hc` is a template parameter so the per-thread row / tree buffers stay in +// registers; see the PERFORMANCE NOTE above. The arithmetic is identical to the +// runtime-`hc` version this replaced. The kernel lives inside the anonymous +// namespace because a template cannot have C linkage; only the dispatcher below +// is `extern "C"`, which is all the Rust FFI binds to. +template __global__ void sinkhorn_normalize_f32_kernel( const float* __restrict__ in, float* __restrict__ out, int n, - int hc, int iters, float eps ) { int batch = blockIdx.x; int row = threadIdx.x; - if (batch >= n || row >= hc) return; + if (batch >= n || row >= HC) return; - extern __shared__ float smem[]; - float* mat = smem; // [hc * hc] - float* csum = smem + hc * hc; // [hc] + __shared__ float mat[HC * HC]; + __shared__ float csum[HC]; - const float* my_in = in + (size_t)batch * hc * hc + (size_t)row * hc; + const float* my_in = in + (size_t)batch * HC * HC + (size_t)row * HC; // Each thread owns one row in registers. - float r[SINKHORN_MAX_HC]; - for (int j = 0; j < hc; ++j) r[j] = my_in[j]; + float r[HC]; +#pragma unroll + for (int j = 0; j < HC; ++j) r[j] = my_in[j]; // ---- 1. stable row softmax, then + eps ---- // candle: max_keepdim(-1) -> broadcast_sub -> exp -> sum_keepdim(-1) // -> broadcast_div -> affine(+eps) - const float m = candle_tree_max(r, hc); - for (int j = 0; j < hc; ++j) r[j] = expf(__fsub_rn(r[j], m)); - const float rs = candle_tree_sum(r, hc); - for (int j = 0; j < hc; ++j) r[j] = __fadd_rn(__fdiv_rn(r[j], rs), eps); + const float m = candle_tree_max(r); +#pragma unroll + for (int j = 0; j < HC; ++j) r[j] = expf(__fsub_rn(r[j], m)); + const float rs = candle_tree_sum(r); +#pragma unroll + for (int j = 0; j < HC; ++j) r[j] = __fadd_rn(__fdiv_rn(r[j], rs), eps); // publish row to shared - for (int j = 0; j < hc; ++j) mat[row * hc + j] = r[j]; +#pragma unroll + for (int j = 0; j < HC; ++j) mat[row * HC + j] = r[j]; __syncthreads(); // ---- 2. initial column normalize: x / (colsum + eps) ---- // column `row` sum = tree-sum over rows k of mat[k][row] { - float col[SINKHORN_MAX_HC]; - for (int k = 0; k < hc; ++k) col[k] = mat[k * hc + row]; - csum[row] = __fadd_rn(candle_tree_sum(col, hc), eps); + float col[HC]; +#pragma unroll + for (int k = 0; k < HC; ++k) col[k] = mat[k * HC + row]; + csum[row] = __fadd_rn(candle_tree_sum(col), eps); } __syncthreads(); - for (int j = 0; j < hc; ++j) r[j] = __fdiv_rn(mat[row * hc + j], csum[j]); +#pragma unroll + for (int j = 0; j < HC; ++j) r[j] = __fdiv_rn(mat[row * HC + j], csum[j]); // ---- 3. (iters - 1) more row->col passes ---- for (int it = 0; it < iters - 1; ++it) { // row normalize: x / (rowsum + eps) (r holds this thread's row) - const float rsum = __fadd_rn(candle_tree_sum(r, hc), eps); - for (int j = 0; j < hc; ++j) r[j] = __fdiv_rn(r[j], rsum); - for (int j = 0; j < hc; ++j) mat[row * hc + j] = r[j]; + const float rsum = __fadd_rn(candle_tree_sum(r), eps); +#pragma unroll + for (int j = 0; j < HC; ++j) r[j] = __fdiv_rn(r[j], rsum); +#pragma unroll + for (int j = 0; j < HC; ++j) mat[row * HC + j] = r[j]; __syncthreads(); // column normalize: x / (colsum + eps) { - float col[SINKHORN_MAX_HC]; - for (int k = 0; k < hc; ++k) col[k] = mat[k * hc + row]; - csum[row] = __fadd_rn(candle_tree_sum(col, hc), eps); + float col[HC]; +#pragma unroll + for (int k = 0; k < HC; ++k) col[k] = mat[k * HC + row]; + csum[row] = __fadd_rn(candle_tree_sum(col), eps); } __syncthreads(); - for (int j = 0; j < hc; ++j) r[j] = __fdiv_rn(mat[row * hc + j], csum[j]); +#pragma unroll + for (int j = 0; j < HC; ++j) r[j] = __fdiv_rn(mat[row * HC + j], csum[j]); } // ---- write out ---- - float* my_out = out + (size_t)batch * hc * hc + (size_t)row * hc; - for (int j = 0; j < hc; ++j) my_out[j] = r[j]; + float* my_out = out + (size_t)batch * HC * HC + (size_t)row * HC; +#pragma unroll + for (int j = 0; j < HC; ++j) my_out[j] = r[j]; } +// Every hc in [1, SINKHORN_MAX_HC] gets its own instantiation, so there is no +// runtime-`hc` fallback whose numerics could drift from the templated path. +// V4-Flash uses hc = 4; the rest are cheap (a few hundred bytes of cubin each) +// and keep the Rust-side contract "hc <= 16 is supported" honest. +#define SINKHORN_LAUNCH(HC_) \ + case HC_: \ + sinkhorn_normalize_f32_kernel<<>>( \ + (const float*)in, (float*)out, n, iters, eps); \ + return; + +} // namespace + +extern "C" { + void sinkhorn_normalize_f32( const void* in, void* out, @@ -198,9 +270,30 @@ void sinkhorn_normalize_f32( ) { dim3 grid(n, 1, 1); dim3 block(hc, 1, 1); - size_t shmem = (size_t)(hc * hc + hc) * sizeof(float); - sinkhorn_normalize_f32_kernel<<>>( - (const float*)in, (float*)out, n, hc, iters, eps); + switch (hc) { + SINKHORN_LAUNCH(1) + SINKHORN_LAUNCH(2) + SINKHORN_LAUNCH(3) + SINKHORN_LAUNCH(4) + SINKHORN_LAUNCH(5) + SINKHORN_LAUNCH(6) + SINKHORN_LAUNCH(7) + SINKHORN_LAUNCH(8) + SINKHORN_LAUNCH(9) + SINKHORN_LAUNCH(10) + SINKHORN_LAUNCH(11) + SINKHORN_LAUNCH(12) + SINKHORN_LAUNCH(13) + SINKHORN_LAUNCH(14) + SINKHORN_LAUNCH(15) + SINKHORN_LAUNCH(16) + default: + // Unreachable: sinkhorn_normalize_cuda rejects hc > SINKHORN_MAX_HC + // before calling. Leaving `out` untouched here would be a silent wrong + // answer, so do nothing and let the caller's guard be the contract. + return; + } } +#undef SINKHORN_LAUNCH } // extern "C" diff --git a/mistralrs-core/src/models/deepseek4.rs b/mistralrs-core/src/models/deepseek4.rs index e03103a1f..f18e6a093 100644 --- a/mistralrs-core/src/models/deepseek4.rs +++ b/mistralrs-core/src/models/deepseek4.rs @@ -2085,11 +2085,23 @@ impl MoeGate { // V4: sqrt(softplus(x)). Stable formulation: // softplus(x) = max(x, 0) + log(1 + exp(-|x|)). // Audit §8 P1 item 14. + // + // The eager form below is NINE kernel launches (`zeros_like`, + // `bmaximum`, `uabs`, `uneg`, `uexp`, `affine`, `ulog`, `badd`, + // `usqrt`) on a `[1, n_routed_experts]` = [1, 256] tensor, once per + // MoE layer per token. `cuda/hc_fused.cu` collapses it to one, + // bit-identically — this expression decides WHICH EXPERTS RUN, so + // the fused kernel transcribes candle's ops rather than + // re-deriving them. `ARC_HC_FUSED=0` restores the chain for A/B. ScoringFunc::SqrtSoftplus => { - let max0 = logits.maximum(&logits.zeros_like()?)?; - let abs = logits.abs()?; - let softplus = (max0 + ((abs.neg()?.exp()? + 1.0)?.log()?))?; - softplus.sqrt()? + if crate::cuda::hc_fused::usable(&logits) { + crate::cuda::hc_fused::sqrt_softplus_cuda(&logits)? + } else { + let max0 = logits.maximum(&logits.zeros_like()?)?; + let abs = logits.abs()?; + let softplus = (max0 + ((abs.neg()?.exp()? + 1.0)?.log()?))?; + softplus.sqrt()? + } } }; drop(_prof_score); diff --git a/mistralrs-core/src/models/dsv4_mhc.rs b/mistralrs-core/src/models/dsv4_mhc.rs index 00afc5940..d3c740ed7 100644 --- a/mistralrs-core/src/models/dsv4_mhc.rs +++ b/mistralrs-core/src/models/dsv4_mhc.rs @@ -278,57 +278,109 @@ impl V4MHCLayerParams { // Promote to F32, flatten to [N, hc*h]. let x_flat = x.reshape((n, hc * h))?.to_dtype(DType::F32)?; - // rsqrt(mean(x^2) + eps) - let sq_mean = x_flat.sqr()?.mean_keepdim(D::Minus1)?; - let rsqrt = (sq_mean + self.rt.rms_norm_eps)?.recip()?.sqrt()?; // [N, 1] - - // mixes = (x_flat @ fn^T) * rsqrt → [N, mix_hc] // Defensively cast weight tensors to F32 — try_load already produces F32, // but hand-constructed callers (tests, external integrators) may not. let hc_fn_f32 = hc_fn.to_dtype(DType::F32)?; let hc_scale_f32 = hc_scale.to_dtype(DType::F32)?; let hc_base_f32 = hc_base.to_dtype(DType::F32)?; let mixes_raw = x_flat.matmul(&hc_fn_f32.t()?)?; - let mixes = mixes_raw.broadcast_mul(&rsqrt)?; - // Slot indices in `mixes`: - // pre : [.., 0 .. hc) - // post : [.., hc .. 2*hc) - // comb : [.., 2*hc .. (2+hc)*hc) reshape to [.., hc, hc] - let pre_block = mixes.narrow(D::Minus1, 0, hc)?; - let post_block = mixes.narrow(D::Minus1, hc, hc)?; - let comb_block = mixes - .narrow(D::Minus1, 2 * hc, hc * hc)? - .reshape((n, hc, hc))?; - - // hc_scale is [3]; hc_base is [mix_hc] split into three blocks. - let s_pre = hc_scale_f32.narrow(0, 0, 1)?; - let s_post = hc_scale_f32.narrow(0, 1, 1)?; - let s_comb = hc_scale_f32.narrow(0, 2, 1)?; - let b_pre = hc_base_f32.narrow(0, 0, hc)?; - let b_post = hc_base_f32.narrow(0, hc, hc)?; - let b_comb = hc_base_f32.narrow(0, 2 * hc, hc * hc)?.reshape((hc, hc))?; - - // pre = sigmoid(pre_block * s_pre + b_pre) + eps - let pre = candle_nn::ops::sigmoid( - &(pre_block.broadcast_mul(&s_pre)?.broadcast_add(&b_pre)?), - )?; - let pre = (pre + self.rt.hc_eps)?; - - // post = 2 * sigmoid(post_block * s_post + b_post) - // NOTE: use affine() for the scalar *2 rather than a device-scalar - // Tensor::new(2f32, device) — the latter is a per-call CPU->GPU sync - // (CLAUDE.md pitfall #5) that breaks CUDA-graph capture of the decode - // forward. affine folds the constant into the kernel, no allocation. - let post_sig = candle_nn::ops::sigmoid( - &(post_block.broadcast_mul(&s_post)?.broadcast_add(&b_post)?), - )?; - let post = post_sig.affine(2.0, 0.0)?; + // Everything from here to `comb_pre` is ONE fused kernel on CUDA. The + // eager chain below it is 18 launches — a 7-launch hand-decomposed RMS + // statistic (`sqr -> fast_sum -> affine -> affine -> recip -> sqrt -> + // bmul`) plus 11 for the three scoring blocks — all on 24 floats once + // the reduction is done, twice per layer, 43 layers. At b=1 that is + // ~1,460 of the step's ~7,900 kernel launches, i.e. pure launch + // overhead. `cuda/hc_fused.cu` is bit-identical to this chain by + // construction, not by tolerance; the eager path stays reachable via + // `ARC_HC_FUSED=0` so the two can be A/B'd from one binary. + let on_cuda = crate::cuda::hc_fused::usable(&x_flat); + let shapes_ok = x_flat.is_contiguous() + && mixes_raw.is_contiguous() + && hc_scale_f32.is_contiguous() + && hc_base_f32.is_contiguous() + && hc_base_f32.dims1().map(|v| v == (2 + hc) * hc).unwrap_or(false) + && hc_scale_f32.dims1().map(|v| v == 3).unwrap_or(false); + if on_cuda && !shapes_ok { + // Falling back on CUDA is a silent 18-launch regression that looks + // exactly like "the fusion didn't help". Say so once rather than + // letting a layout change quietly undo the optimisation. + static WARNED: std::sync::Once = std::sync::Once::new(); + WARNED.call_once(|| { + tracing::warn!( + "V4 mHC: fused hc_pre kernel unusable (x_flat contig={}, mixes contig={}, \ + scale contig={} dims={:?}, base contig={} dims={:?}) — falling back to the \ + 18-launch eager chain.", + x_flat.is_contiguous(), + mixes_raw.is_contiguous(), + hc_scale_f32.is_contiguous(), + hc_scale_f32.dims(), + hc_base_f32.is_contiguous(), + hc_base_f32.dims(), + ); + }); + } + let fused = on_cuda && shapes_ok; + + let (pre, post, comb_pre) = if fused { + crate::cuda::hc_fused::hc_pre_fused_cuda( + &x_flat, + &mixes_raw, + &hc_scale_f32, + &hc_base_f32, + hc, + self.rt.rms_norm_eps, + self.rt.hc_eps, + )? + } else { + // rsqrt(mean(x^2) + eps) + let sq_mean = x_flat.sqr()?.mean_keepdim(D::Minus1)?; + let rsqrt = (sq_mean + self.rt.rms_norm_eps)?.recip()?.sqrt()?; // [N, 1] + + // mixes = (x_flat @ fn^T) * rsqrt → [N, mix_hc] + let mixes = mixes_raw.broadcast_mul(&rsqrt)?; + + // Slot indices in `mixes`: + // pre : [.., 0 .. hc) + // post : [.., hc .. 2*hc) + // comb : [.., 2*hc .. (2+hc)*hc) reshape to [.., hc, hc] + let pre_block = mixes.narrow(D::Minus1, 0, hc)?; + let post_block = mixes.narrow(D::Minus1, hc, hc)?; + let comb_block = mixes + .narrow(D::Minus1, 2 * hc, hc * hc)? + .reshape((n, hc, hc))?; + + // hc_scale is [3]; hc_base is [mix_hc] split into three blocks. + let s_pre = hc_scale_f32.narrow(0, 0, 1)?; + let s_post = hc_scale_f32.narrow(0, 1, 1)?; + let s_comb = hc_scale_f32.narrow(0, 2, 1)?; + let b_pre = hc_base_f32.narrow(0, 0, hc)?; + let b_post = hc_base_f32.narrow(0, hc, hc)?; + let b_comb = hc_base_f32.narrow(0, 2 * hc, hc * hc)?.reshape((hc, hc))?; + + // pre = sigmoid(pre_block * s_pre + b_pre) + eps + let pre = candle_nn::ops::sigmoid( + &(pre_block.broadcast_mul(&s_pre)?.broadcast_add(&b_pre)?), + )?; + let pre = (pre + self.rt.hc_eps)?; + + // post = 2 * sigmoid(post_block * s_post + b_post) + // NOTE: use affine() for the scalar *2 rather than a device-scalar + // Tensor::new(2f32, device) — the latter is a per-call CPU->GPU sync + // (CLAUDE.md pitfall #5) that breaks CUDA-graph capture of the decode + // forward. affine folds the constant into the kernel, no allocation. + let post_sig = candle_nn::ops::sigmoid( + &(post_block.broadcast_mul(&s_post)?.broadcast_add(&b_post)?), + )?; + let post = post_sig.affine(2.0, 0.0)?; + + let comb_pre = comb_block + .broadcast_mul(&s_comb)? + .broadcast_add(&b_comb)?; // [N, hc, hc] + (pre, post, comb_pre) + }; // comb = sinkhorn_normalize(comb_block * s_comb + b_comb) - let comb_pre = comb_block - .broadcast_mul(&s_comb)? - .broadcast_add(&b_comb)?; // [N, hc, hc] let comb = sinkhorn_normalize(&comb_pre, self.rt.hc_sinkhorn_iters, self.rt.hc_eps)?; // y = sum_i pre[..., i, None] * x[..., i, :] → [N, hidden] From a92cf6d11f720480e38bc5a47fe0f007fb93cb58 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick Date: Wed, 19 Aug 2026 22:05:45 +0100 Subject: [PATCH 20/27] =?UTF-8?q?fix(ArcMoE):=20cuda-gated=20build=20?= =?UTF-8?q?=E2=80=94=20device=5Fptr=20guard=20lifetime=20and=20monomorphic?= =?UTF-8?q?=20wrap=20closure?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SpVNMpb13HkUXqSqbN1o9H --- mistralrs-core/src/cuda/hc_fused.rs | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/mistralrs-core/src/cuda/hc_fused.rs b/mistralrs-core/src/cuda/hc_fused.rs index c176e93ea..e0ad69425 100644 --- a/mistralrs-core/src/cuda/hc_fused.rs +++ b/mistralrs-core/src/cuda/hc_fused.rs @@ -60,7 +60,12 @@ mod cuda_impl { candle::Storage::Cuda(c) => c.as_cuda_slice::()?, _ => candle::bail!("hc_fused: {what} must be on CUDA"), }; - Ok(s.slice(l.start_offset()..).device_ptr(s.stream()).0 as *const std::ffi::c_void) + // Bind before returning: `device_ptr` hands back a (ptr, guard) pair + // whose guard borrows `s`, so the pointer must be extracted in its own + // statement rather than in tail position. Same shape as + // `sinkhorn::sinkhorn_normalize_cuda`. + let ptr = s.slice(l.start_offset()..).device_ptr(s.stream()).0 as *const std::ffi::c_void; + Ok(ptr) } /// Fused `hc_pre` middle section: the RMS statistic, its broadcast into @@ -150,14 +155,13 @@ mod cuda_impl { ); } - let wrap = |buf, shape| { - let st = candle::CudaStorage::wrap_cuda_slice(buf, dev.clone()); - Tensor::from((candle::Storage::Cuda(st), shape)) - }; + let pre_st = candle::CudaStorage::wrap_cuda_slice(pre_buf, dev.clone()); + let post_st = candle::CudaStorage::wrap_cuda_slice(post_buf, dev.clone()); + let comb_st = candle::CudaStorage::wrap_cuda_slice(comb_buf, dev.clone()); Ok(( - wrap(pre_buf, (n, hc)), - wrap(post_buf, (n, hc)), - wrap(comb_buf, (n, hc, hc)), + Tensor::from((candle::Storage::Cuda(pre_st), (n, hc))), + Tensor::from((candle::Storage::Cuda(post_st), (n, hc))), + Tensor::from((candle::Storage::Cuda(comb_st), (n, hc, hc))), )) } From 48ae6b70ca9aa66404ec2ba6d3846a54350a4825 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick Date: Wed, 19 Aug 2026 22:13:16 +0100 Subject: [PATCH 21/27] test(ArcMoE): extend the build-wiring tripwire to hc_fused.cu MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The sinkhorn guard asserted build.rs contained the literal `.exclude(&["sinkhorn.cu"])`, so adding a second bit-identity-critical kernel to that list failed it for the one reason that is not a regression — the kind of failure that gets 'fixed' by deleting the assertion. It now matches on the exclude list's CONTENTS, and hc_fused.rs carries the mirror guard: fast-math #error present, IEEE intrinsics present, __expf/__logf/ __fdividef/rsqrtf absent, and the file wired into the --fmad=false builder. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SpVNMpb13HkUXqSqbN1o9H --- mistralrs-core/src/cuda/hc_fused.rs | 46 +++++++++++++++++++++++++++++ mistralrs-core/src/cuda/sinkhorn.rs | 18 +++++++++-- 2 files changed, 62 insertions(+), 2 deletions(-) diff --git a/mistralrs-core/src/cuda/hc_fused.rs b/mistralrs-core/src/cuda/hc_fused.rs index e0ad69425..7376ce005 100644 --- a/mistralrs-core/src/cuda/hc_fused.rs +++ b/mistralrs-core/src/cuda/hc_fused.rs @@ -444,6 +444,52 @@ mod tests { } } + /// Source-level tripwires, mirroring `sinkhorn::tests`: the IEEE + /// intrinsics and the fast-math `#error` guard must stay in hc_fused.cu, + /// fast-math approximations must stay out, and build.rs must keep the file + /// out of the `--use_fast_math` glob and in the `--fmad=false` builder. + /// Getting this wiring wrong is silent — the kernel still runs, it just + /// stops being bit-identical — so the `#error` guard is the hard stop and + /// these string checks catch it on CPU CI too. + #[test] + fn kernel_source_and_build_wiring_guards() { + let cu = include_str!("hc_fused.cu"); + assert!( + cu.contains("#if defined(__USE_FAST_MATH__)") && cu.contains("#error"), + "hc_fused.cu lost its fast-math #error guard" + ); + for required in ["__fadd_rn", "__fmul_rn", "candle_recip", "candle_sigmoid"] { + assert!(cu.contains(required), "hc_fused.cu lost required token {required}"); + } + for forbidden in ["__expf(", "__logf(", "__fdividef(", "rsqrtf(", "__frsqrt_rn("] { + assert!( + !cu.contains(forbidden), + "hc_fused.cu contains {forbidden}, which is not what candle-kernels computes" + ); + } + + let build = include_str!("../../build.rs"); + let exclude = build + .split(".exclude(&[") + .nth(1) + .and_then(|s| s.split("])").next()) + .expect("build.rs no longer calls .exclude(&[..]) on the fast-math builder"); + assert!( + exclude.contains("\"hc_fused.cu\""), + "build.rs no longer excludes hc_fused.cu from the fast-math builder \ + (exclude list is: {exclude})" + ); + assert!( + build.contains(r#""src/cuda/hc_fused.cu""#), + "build.rs no longer feeds hc_fused.cu to the IEEE (no-fast-math) builder" + ); + assert!( + build.contains("--fmad=false"), + "build.rs lost --fmad=false, so nvcc may contract mul+add into an FMA and break \ + bit-identity with candle's unfused op chain" + ); + } + #[test] fn candle_reduce_block_dim_matches_fast_reduce() { // cuda_backend/mod.rs: usize::min(1024, el_to_sum_per_block).next_power_of_two() diff --git a/mistralrs-core/src/cuda/sinkhorn.rs b/mistralrs-core/src/cuda/sinkhorn.rs index 40173853b..c101c7d4c 100644 --- a/mistralrs-core/src/cuda/sinkhorn.rs +++ b/mistralrs-core/src/cuda/sinkhorn.rs @@ -474,9 +474,23 @@ mod tests { } let build = include_str!("../../build.rs"); + // The exclude list grows as more bit-identity-critical kernels are + // added (hc_fused.cu joined it), so match on the list's CONTENTS rather + // than on its exact spelling — otherwise this guard fails for the one + // reason that is not a regression, and gets "fixed" by weakening it. + let exclude = build + .split(".exclude(&[") + .nth(1) + .and_then(|s| s.split("])").next()) + .expect("build.rs no longer calls .exclude(&[..]) on the fast-math builder"); assert!( - build.contains(r#".exclude(&["sinkhorn.cu"])"#), - "build.rs no longer excludes sinkhorn.cu from the fast-math builder" + exclude.contains("\"sinkhorn.cu\""), + "build.rs no longer excludes sinkhorn.cu from the fast-math builder \ + (exclude list is: {exclude})" + ); + assert!( + build.contains(r#""src/cuda/sinkhorn.cu""#), + "build.rs no longer feeds sinkhorn.cu to the IEEE (no-fast-math) builder" ); assert!( build.contains("--fmad=false"), From ab5fbd54d38b228fe0438b2b2ddbae95567016a3 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick Date: Wed, 19 Aug 2026 22:32:41 +0100 Subject: [PATCH 22/27] docs(ArcMoE): cite the measured H200 numbers, not the estimates Measured on H200, git 6f4cd0dd, bin f6ec4e85, toggled with ARC_HC_FUSED on one binary so the comparison has no second variable: ARC_HC_FUSED=0 7,494.3 launches/step 10,892.4 allocs/step 17.3 tok/s 57.70 ms/T ARC_HC_FUSED=1 5,691.0 launches/step 8,918.5 allocs/step 20.1 tok/s 49.66 ms/T The off leg reproduces origin/master (7,477.9 launches, 17.0 tok/s) to within 0.2%, so the switch is the only thing that moved and nothing else regressed. Router region 76 -> 38 kernels/layer, 0.207 -> 0.100 ms GPU. Sinkhorn 30.03 -> 9.12 us/call (2.583 -> 0.785 ms/step). Output bit-identical across the toggle. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SpVNMpb13HkUXqSqbN1o9H --- mistralrs-core/src/cuda/hc_fused.cu | 22 +++++++++++++++------- mistralrs-core/src/models/dsv4_mhc.rs | 8 +++++--- 2 files changed, 20 insertions(+), 10 deletions(-) diff --git a/mistralrs-core/src/cuda/hc_fused.cu b/mistralrs-core/src/cuda/hc_fused.cu index bdfc711ce..a189822aa 100644 --- a/mistralrs-core/src/cuda/hc_fused.cu +++ b/mistralrs-core/src/cuda/hc_fused.cu @@ -5,13 +5,21 @@ // --------------------------------------------------------------------------- // WHY // --------------------------------------------------------------------------- -// At b=1 the V4 decode step is launch-bound, not bandwidth-bound: the measured -// step issues ~7.9k `cuLaunchKernel` calls, of which the router region alone is -// 3,259 (41.8%) spread over 86 contiguous spans -- exactly two per layer -// (`mhc_attn_pre` and `mhc_ffn_pre`, 43 layers). Those spans operate on -// `[1, 24]` / `[1, 4, 4]` / `[1, 256]` tensors: mean kernel duration 2.79 us, -// 89% of them under 5 us. They are pure launch overhead wearing a kernel -// costume. +// At b=1 the V4 decode step is launch-bound, not bandwidth-bound. Measured on +// an H200 at master ef581e9c8: 7,477.9 `cuLaunchKernel` calls per decode step +// costing 18.976 ms/step of host time, of which the router region is 76 kernels +// per layer (3,268/step, 43.7%) spread over 86 contiguous spans -- exactly two +// per layer (`mhc_attn_pre` and `mhc_ffn_pre`, 43 layers). Those spans operate +// on `[1, 24]` / `[1, 4, 4]` / `[1, 256]` tensors: mean kernel duration 2.72 us, +// 89% of them under 5 us, 28.4% of all kernel time. They are pure launch +// overhead wearing a kernel costume. +// +// MEASURED RESULT of this file plus the sinkhorn fix, same box, same binary, +// toggled with ARC_HC_FUSED: 7,494.3 -> 5,691.0 launches/step (-24.1%), +// 10,892.4 -> 8,918.5 allocations/step, 57.70 -> 49.66 ms/token +// (17.3 -> 20.1 tok/s). The router region itself goes 76 -> 38 kernels/layer +// and 0.207 -> 0.100 ms of GPU time. Generated tokens and logprobs are +// bit-identical across the toggle (6 prompts, 768 logprob values, 0 mismatches). // // Two expressions dominate the count and are collapsed here: // diff --git a/mistralrs-core/src/models/dsv4_mhc.rs b/mistralrs-core/src/models/dsv4_mhc.rs index d3c740ed7..57fb1e2d7 100644 --- a/mistralrs-core/src/models/dsv4_mhc.rs +++ b/mistralrs-core/src/models/dsv4_mhc.rs @@ -290,10 +290,12 @@ impl V4MHCLayerParams { // statistic (`sqr -> fast_sum -> affine -> affine -> recip -> sqrt -> // bmul`) plus 11 for the three scoring blocks — all on 24 floats once // the reduction is done, twice per layer, 43 layers. At b=1 that is - // ~1,460 of the step's ~7,900 kernel launches, i.e. pure launch - // overhead. `cuda/hc_fused.cu` is bit-identical to this chain by + // ~1,460 of the decode step's measured 7,494 kernel launches, i.e. pure + // launch overhead. `cuda/hc_fused.cu` is bit-identical to this chain by // construction, not by tolerance; the eager path stays reachable via - // `ARC_HC_FUSED=0` so the two can be A/B'd from one binary. + // `ARC_HC_FUSED=0` so the two can be A/B'd from one binary — and they + // were: flipping it moves 1,803 launches/step and 8.04 ms/token while + // leaving 6 greedy completions and their 768 logprobs bit-identical. let on_cuda = crate::cuda::hc_fused::usable(&x_flat); let shapes_ok = x_flat.is_contiguous() && mixes_raw.is_contiguous() From 179e405ac29ab4717c504a607cc8fc9669a98246 Mon Sep 17 00:00:00 2001 From: arc-agent Date: Thu, 20 Aug 2026 03:45:43 +0000 Subject: [PATCH 23/27] =?UTF-8?q?perf(ArcInfer/ArcMoE):=20fuse=20the=20mHC?= =?UTF-8?q?=20epilogue=20and=20post-step=20=E2=80=94=2014=20launches=20per?= =?UTF-8?q?=20mixing=20point=20become=207?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit At b=1 the V4 decode step is launch-bound, not bandwidth-bound. After the router-region fusion (8b79311) one mHC mixing point was still FOURTEEN kernel launches, and there are 86 of them per token (43 layers x {attn, ffn}). Two blocks are collapsed here, both operating on [1, 4, 4096] tensors: hc_pre tail `pre.unsqueeze(-1).broadcast_mul(x_f32).sum(1)` + the narrowing cast — 3 launches to compute a 4-term weighted average. hc_post whole — 6 launches whose real work is 4 multiply-adds per output. MEASURED, H200, clocks locked at 1980 MHz, ONE binary toggled with ARC_HC_FUSE2 (a two-binary before/after would confound the change with the rebuild): launches/step 6,192.5 -> 5,566.4 (-626.1, -10.1%) ms/token 40.50 -> 38.18 (24.7 -> 26.2 tok/s, +6.1%) GPU time/step -1.48 ms Every kernel that moved is one of the nine this touches, each at 86/step: cast_bf16_f32 -172.2, cast_f32_bf16 -172.1, bmul_f32 -172.0, fast_sum_f32 -86.1, badd_f32 -86.0, sm80_xmma_gemm -86.0, against hc_post_fused +85.9 and hc_y_combine +85.9. Nothing else moved: arc_kv_fp8_quantize 42.98 -> 42.95, qtip_gather_gemv 128.93 -> 128.86, dot_kernel 86.95 -> 86.90. The gate GEMV is deliberately left alone. cuBLAS splits the [1,16384]x[16384,24] product into dot_kernel + reduce_1Block, and its K=16384 reduction order is not reproducible by construction. Fusing it is the only way to merge hc_post with the next hc_pre into one cross-layer kernel; that trade — one launch per seam against a tolerance-based result — is not taken. BIT-IDENTICAL, and proven on live data rather than argued. ARC_HC_AB=1 compared 4,000 tensors against the eager chain with zero mismatching bits. ARC_HC_AB_POISON=1 perturbs one element by a single ULP and the same comparison flags all 4,000 (fused=0xbe32 eager=0xbe31), so the check is live rather than vacuous. hc_fused.cu has cited `ARC_HC_AB=1` as "the final proof" since it was written; the string appeared nowhere else in the tree. This adds it. Also records a measured defect rather than fixing it silently: the `#if defined(__USE_FAST_MATH__)` guard in BOTH hc_fused.cu and sinkhorn.cu is dead. nvcc 12.4 defines neither __USE_FAST_MATH__ nor __FAST_MATH__ in either pass — nvcc --use_fast_math -E -dM x.cu | grep -i fast -> no output nvcc --use_fast_math -arch=sm_90 -c guard.cu -> compiles clean — while --use_fast_math does reach the device pass (a float divide drops from 3 MUFU/RCP instructions to 1). Two bit-identity-critical files were relying on a macro that cannot exist. The real protection is the build.rs wiring already asserted in tests plus the runtime A/B added here. Co-Authored-By: Claude Opus 5 (1M context) --- mistralrs-core/src/cuda/ffi.rs | 28 ++ mistralrs-core/src/cuda/hc_fused.cu | 302 +++++++++++++ mistralrs-core/src/cuda/hc_fused.rs | 588 +++++++++++++++++++++++++- mistralrs-core/src/models/dsv4_mhc.rs | 111 +++-- 4 files changed, 996 insertions(+), 33 deletions(-) diff --git a/mistralrs-core/src/cuda/ffi.rs b/mistralrs-core/src/cuda/ffi.rs index a2ca78e45..441e0d954 100644 --- a/mistralrs-core/src/cuda/ffi.rs +++ b/mistralrs-core/src/cuda/ffi.rs @@ -399,4 +399,32 @@ extern "C" { numel: i64, stream: i64, ); + /// `y = sum_i pre[i] * x_f32[i, :]` narrowed to the model dtype. + /// Returns 0 on launch, non-zero if `(hc, dtype)` has no specialisation — + /// the caller must then use the eager chain. + #[allow(clippy::too_many_arguments)] + pub(crate) fn hc_y_combine( + x_f32: *const c_void, + pre: *const c_void, + y: *mut c_void, + hc: i32, + h: i32, + total: i64, + dtype: i32, + stream: i64, + ) -> i32; + /// The whole `hc_post` re-expansion. Same return contract as above. + #[allow(clippy::too_many_arguments)] + pub(crate) fn hc_post_fused( + x: *const c_void, + residual: *const c_void, + post: *const c_void, + comb: *const c_void, + out: *mut c_void, + hc: i32, + h: i32, + total: i64, + dtype: i32, + stream: i64, + ) -> i32; } diff --git a/mistralrs-core/src/cuda/hc_fused.cu b/mistralrs-core/src/cuda/hc_fused.cu index a189822aa..2ff941e57 100644 --- a/mistralrs-core/src/cuda/hc_fused.cu +++ b/mistralrs-core/src/cuda/hc_fused.cu @@ -87,8 +87,31 @@ // alongside the fused kernel and reports any bitwise mismatch. #include +#include #include +// --------------------------------------------------------------------------- +// THE COMPILE-TIME FAST-MATH GUARD BELOW IS DEAD. MEASURED, NOT ASSUMED. +// --------------------------------------------------------------------------- +// `nvcc --use_fast_math` defines NEITHER `__USE_FAST_MATH__` NOR `__FAST_MATH__` +// in either the host or the device pass. Probed on this box, CUDA 12.4: +// +// $ nvcc --use_fast_math -E -dM x.cu | grep -i fast -> (no output) +// $ nvcc --use_fast_math -arch=sm_90 -c guard.cu -> compiles clean +// (with #error on both macros present) +// +// So the `#if defined(__USE_FAST_MATH__)` below has never once been capable of +// firing -- in this file or in sinkhorn.cu, which carries the same guard. It is +// left in place ONLY so that this comment sits next to it; do not read it as +// protection. The hazard it was written for is real: the same probe shows +// `--use_fast_math` cutting a float divide from 3 MUFU/RCP instructions to 1, +// i.e. it does reach the device pass and does change the arithmetic. +// +// The guard that actually works is at RUNTIME and on REAL DATA: `ARC_HC_AB=1` +// recomputes the eager candle chain next to every fused kernel and reports the +// first bitwise mismatch (see `ab_check` in hc_fused.rs). It subsumes the +// compile-flag question -- a fast-math build fails it on the first token -- and +// it has been proven RED by injecting a 1-ULP perturbation. #if defined(__USE_FAST_MATH__) #error "hc_fused.cu must be compiled WITHOUT --use_fast_math: fast math rewrites expf/logf to the hardware approximations and IEEE division to approximate reciprocals, breaking bit-identity with candle-kernels (which build with plain -O3). See the dedicated no-fast-math builder in mistralrs-core/build.rs." #endif @@ -297,3 +320,282 @@ void sqrt_softplus_f32(const void *inp, void *out, long long numel, long long st } } // extern "C" + +// =========================================================================== +// The mHC epilogue and post-step, one launch each. +// =========================================================================== +// +// After the router-region fusion above, one mHC mixing point is still FOURTEEN +// launches at b=1 -- and there are 86 of them per token (43 layers x {attn, +// ffn}). Counted from an nsys trace on an H200, not from reading the source: +// per mixing point, cast_bf16_f32 x3, cast_f32_bf16 x2, bmul_f32 x2, +// fast_sum_f32, badd_f32, sm80_xmma_gemm, dot_kernel + reduce_1Block (cuBLAS +// splits the [1,16384]x[16384,24] gate GEMV into two), hc_pre_fused, sinkhorn. +// +// The two blocks collapsed here are the nine that are not the gate GEMV: +// +// hc_pre tail `pre.unsqueeze(-1).broadcast_mul(x_f32).sum(1)` + narrowing +// -- 3 launches to compute a 4-term weighted average. +// hc_post whole `to_dtype x` + `to_dtype residual` + `broadcast_mul` + +// `transpose.contiguous` + `matmul` + `add` + `to_dtype` +// -- 6 launches whose real work is 4 multiply-adds per output. +// +// MEASURED RESULT, H200, clocks locked at 1980 MHz, ONE binary toggled with +// ARC_HC_FUSE2 (a two-binary before/after would confound this with the +// rebuild): 6,192.5 -> 5,566.4 launches/step (-626.1, -10.1%) and 40.50 -> +// 38.18 ms/token (24.7 -> 26.2 tok/s). Every kernel that moved is one of these +// nine, each at 86/step: cast_bf16_f32 -172.2, cast_f32_bf16 -172.1, +// bmul_f32 -172.0, fast_sum_f32 -86.1, badd_f32 -86.0, sm80_xmma_gemm -86.0, +// against hc_post_fused +85.9 and hc_y_combine +85.9. GPU time falls too, by +// 1.48 ms/step, because the intermediate [1,4,4096] F32 tensors are never +// materialised. +// +// Both operate on [1, 4, 7168] at b=1: ~114 kB of traffic behind ~9 kernel +// launches. They are launch overhead wearing a kernel costume, the same +// diagnosis this file already applied to the scoring block. +// +// The GEMM in `hc_pre` (`x_flat @ hc_fn.T`, [1,28672]x[28672,24]) is +// DELIBERATELY LEFT ALONE. It is cuBLAS, its K=28672 reduction order is not +// reproducible by construction, and bit-identity is the contract here. Fusing +// it would be the only way to merge `hc_post` with the next `hc_pre` into a +// single cross-layer kernel; that trade -- one launch per seam against a +// tolerance-based result -- is not taken. + +namespace { + +// Widen to f32. BF16->F32 is an exact widening, so this introduces no rounding +// and the fused path sees exactly the bits candle's `to_dtype(F32)` produced. +__device__ __forceinline__ float hc_widen(float v) { return v; } +__device__ __forceinline__ float hc_widen(__nv_bfloat16 v) { return __bfloat162float(v); } + +// Narrow back to the model dtype. candle's cast kernel is +// `static_cast<__nv_bfloat16>(x)`, i.e. `__float2bfloat16`, round-to-nearest- +// even -- the same intrinsic, not an approximation of it. +template __device__ __forceinline__ T hc_narrow(float v); +template <> __device__ __forceinline__ float hc_narrow(float v) { return v; } +template <> __device__ __forceinline__ __nv_bfloat16 hc_narrow<__nv_bfloat16>(float v) { + return __float2bfloat16(v); +} + +// --------------------------------------------------------------------------- +// y = sum_i pre[i] * x[i, :] (hc_pre tail, 3 launches -> 1) +// --------------------------------------------------------------------------- +// +// Replaces verbatim (dsv4_mhc.rs `hc_pre`): +// let y = pre.unsqueeze(-1)?.broadcast_mul(&x_f32)?.sum(1)?; // 2 +// let y_out = y.reshape(y_shape)?.to_dtype(in_dtype)?; // 1 +// +// REDUCTION ORDER -- read from candle, not remembered. `sum(1)` on [n, hc, h] +// moves the summed dim last (cuda_backend/mod.rs `FastReduce`), giving +// dst_el = n*h and el_to_sum_per_block = hc, hence +// block_dim = min(1024, hc).next_power_of_two() = hc for a power-of-two hc. +// candle-kernels `fast_sum` (reduce.cu:70-103) then runs +// +// shr[tid] = 0; // note: canonicalises -0.0f to +0.0f +// while (idx < stop) shr[tid] += src[strided_i]; // ONE element per thread +// for (s = bd/2; s; s >>= 1) { __syncthreads(); if (tid +__global__ void hc_y_combine_kernel( + const float *__restrict__ x_f32, // [n, hc, h] F32 contiguous + const float *__restrict__ pre, // [n, hc] F32 + OUT_T *__restrict__ y, // [n, h] out, model dtype + int h, + long long total // n * h +) { + const long long stride = (long long)blockDim.x * (long long)gridDim.x; + for (long long i = (long long)blockIdx.x * (long long)blockDim.x + threadIdx.x; + i < total; i += stride) { + const long long row = i / (long long)h; + const int col = (int)(i - row * (long long)h); + + const float *prow = pre + row * HC; + const float *xrow = x_f32 + row * (long long)HC * (long long)h + col; + + float acc[HC]; +#pragma unroll + for (int t = 0; t < HC; ++t) { + // `shr[tid] = 0` then `shr[tid] += a_t`: the add-to-zero is part of + // candle's kernel and is NOT a no-op for a_t == -0.0f. + acc[t] = __fadd_rn(0.0f, __fmul_rn(prow[t], xrow[(long long)t * (long long)h])); + } +#pragma unroll + for (int s = HC >> 1; s > 0; s >>= 1) { +#pragma unroll + for (int t = 0; t < HC; ++t) { + if (t < s) { + acc[t] = __fadd_rn(acc[t], acc[t + s]); + } + } + } + y[i] = hc_narrow(acc[0]); + } +} + +// --------------------------------------------------------------------------- +// hc_post: re-expand one branch output into the hc residual streams (6 -> 1) +// --------------------------------------------------------------------------- +// +// Replaces verbatim (dsv4_mhc.rs `hc_post`): +// let x_n = x.reshape((n,h))?.to_dtype(F32)?; // 1 +// let residual_n = residual.reshape((n,hc,h))?.to_dtype(F32)?; // 1 +// let term1 = post_n.unsqueeze(-1)?.broadcast_mul(&x_n.unsqueeze(1)?)?; // 1 +// let term2 = comb_n.transpose(1,2)?.matmul(&residual_n)?; // 2 +// (the transpose forces a contiguous copy) +// let out_n = (term1 + term2)?; // 1 +// out_n.reshape(out_shape)?.to_dtype(in_dtype) // 1 +// +// term2 is a [hc,hc] x [hc,h] product with K = hc = 4. The K-loop is written +// here as the sequential `acc = fmaf(c, r, acc)` from a zero accumulator, which +// is what a cuBLAS GEMM with K=4 does; unlike the K=28672 GEMM in `hc_pre` this +// is short enough that the order is checkable rather than assumed, and it IS +// checked -- `ARC_HC_AB=1` compares this kernel against the eager chain above +// bit for bit on live decode data. +template +__global__ void hc_post_fused_kernel( + const T *__restrict__ x, // [n, h] model dtype + const T *__restrict__ residual, // [n, hc, h] model dtype + const float *__restrict__ post, // [n, hc] F32 + const float *__restrict__ comb, // [n, hc, hc] F32 + T *__restrict__ out, // [n, hc, h] model dtype + int h, + long long total // n * hc * h +) { + const long long stride = (long long)blockDim.x * (long long)gridDim.x; + const long long hh = (long long)HC * (long long)h; + for (long long i = (long long)blockIdx.x * (long long)blockDim.x + threadIdx.x; + i < total; i += stride) { + const long long row = i / hh; + const long long rem = i - row * hh; + const int k = (int)(rem / (long long)h); + const int col = (int)(rem - (long long)k * (long long)h); + + // term1 = post[k] * x[col]; candle's broadcast_mul is binary.cu `x * y`, + // a bare multiply with nothing to contract into an FMA. + const float xv = hc_widen(x[row * (long long)h + col]); + const float t1 = __fmul_rn(post[row * HC + k], xv); + + // term2 = sum_j comb[j][k] * residual[j][col] + const float *crow = comb + row * (long long)HC * (long long)HC; + const T *rrow = residual + row * hh + col; + float acc = 0.0f; +#pragma unroll + for (int j = 0; j < HC; ++j) { + acc = fmaf(crow[j * HC + k], hc_widen(rrow[(long long)j * (long long)h]), acc); + } + + out[i] = hc_narrow(__fadd_rn(t1, acc)); + } +} + +// 256 threads, grid capped so the grid-stride loop stays valid for any size. +inline void hc_launch_dims(long long total, unsigned &blocks, int &block) { + block = 256; + long long b = (total + block - 1) / block; + if (b < 1) { + b = 1; + } + if (b > 65535) { + b = 65535; + } + blocks = (unsigned)b; +} + +} // namespace + +extern "C" { + +// dtype tags shared with hc_fused.rs. Anything else must not reach these. +#define HC_DTYPE_F32 0 +#define HC_DTYPE_BF16 1 + +// Returns 0 on launch, non-zero if (hc, dtype) is outside the specialised set — +// the caller then falls back to the eager chain rather than producing a wrong +// answer with a generic slow path. +int hc_y_combine( + const void *x_f32, + const void *pre, + void *y, + int hc, + int h, + long long total, + int dtype, + long long stream +) { + unsigned blocks; + int block; + hc_launch_dims(total, blocks, block); + cudaStream_t s = (cudaStream_t)stream; +#define HC_YC(T, HCV) \ + hc_y_combine_kernel<<>>( \ + (const float *)x_f32, (const float *)pre, (T *)y, h, total) + if (dtype == HC_DTYPE_F32) { + switch (hc) { + case 2: HC_YC(float, 2); return 0; + case 4: HC_YC(float, 4); return 0; + case 8: HC_YC(float, 8); return 0; + case 16: HC_YC(float, 16); return 0; + default: return 1; + } + } else if (dtype == HC_DTYPE_BF16) { + switch (hc) { + case 2: HC_YC(__nv_bfloat16, 2); return 0; + case 4: HC_YC(__nv_bfloat16, 4); return 0; + case 8: HC_YC(__nv_bfloat16, 8); return 0; + case 16: HC_YC(__nv_bfloat16, 16); return 0; + default: return 1; + } + } + return 1; +#undef HC_YC +} + +int hc_post_fused( + const void *x, + const void *residual, + const void *post, + const void *comb, + void *out, + int hc, + int h, + long long total, + int dtype, + long long stream +) { + unsigned blocks; + int block; + hc_launch_dims(total, blocks, block); + cudaStream_t s = (cudaStream_t)stream; +#define HC_PF(T, HCV) \ + hc_post_fused_kernel<<>>( \ + (const T *)x, (const T *)residual, (const float *)post, (const float *)comb, \ + (T *)out, h, total) + if (dtype == HC_DTYPE_F32) { + switch (hc) { + case 2: HC_PF(float, 2); return 0; + case 4: HC_PF(float, 4); return 0; + case 8: HC_PF(float, 8); return 0; + case 16: HC_PF(float, 16); return 0; + default: return 1; + } + } else if (dtype == HC_DTYPE_BF16) { + switch (hc) { + case 2: HC_PF(__nv_bfloat16, 2); return 0; + case 4: HC_PF(__nv_bfloat16, 4); return 0; + case 8: HC_PF(__nv_bfloat16, 8); return 0; + case 16: HC_PF(__nv_bfloat16, 16); return 0; + default: return 1; + } + } + return 1; +#undef HC_PF +} + +} // extern "C" diff --git a/mistralrs-core/src/cuda/hc_fused.rs b/mistralrs-core/src/cuda/hc_fused.rs index 7376ce005..d893340c8 100644 --- a/mistralrs-core/src/cuda/hc_fused.rs +++ b/mistralrs-core/src/cuda/hc_fused.rs @@ -31,6 +31,19 @@ pub fn fused_enabled() -> bool { *ENABLED.get_or_init(|| !matches!(std::env::var("ARC_HC_FUSED").as_deref(), Ok("0"))) } +/// `ARC_HC_FUSE2=0` disables ONLY the second-wave fusions — the `hc_pre` +/// weighted-average tail and the whole of `hc_post` — while leaving the +/// router-region kernels (`hc_pre_fused`, `sqrt_softplus`) on. +/// +/// Separate from `ARC_HC_FUSED` on purpose: with one switch per wave, the +/// contribution of THIS change is measurable from a single binary, so a +/// before/after cannot be contaminated by a rebuild. +pub fn fuse2_enabled() -> bool { + use std::sync::OnceLock; + static ENABLED: OnceLock = OnceLock::new(); + *ENABLED.get_or_init(|| !matches!(std::env::var("ARC_HC_FUSE2").as_deref(), Ok("0"))) +} + /// candle's `FastReduce` block size for a reduction of `len` elements /// (`cuda_backend/mod.rs`: `usize::min(1024, el_to_sum_per_block) /// .next_power_of_two()`). The fused kernel must use exactly this, because the @@ -195,10 +208,178 @@ mod cuda_impl { logits.shape().clone(), ))) } + + /// Tag matching `HC_DTYPE_*` in hc_fused.cu. Only the dtypes whose cast + /// semantics have been transcribed are accepted; anything else returns an + /// error so the caller falls back to the eager chain rather than silently + /// running a kernel that rounds differently. + fn dtype_tag(d: DType) -> Result { + match d { + DType::F32 => Ok(0), + DType::BF16 => Ok(1), + other => candle::bail!("hc_fused: unsupported dtype {other:?}"), + } + } + + /// Device pointer for a contiguous CUDA tensor of a supported dtype. + fn any_ptr(t: &Tensor, what: &str) -> Result<*const std::ffi::c_void> { + use candle_core::cuda_backend::cudarc::driver::DevicePtr; + if !t.is_contiguous() { + candle::bail!("hc_fused: {what} must be contiguous"); + } + let (s, l) = t.storage_and_layout(); + let cuda = match &*s { + candle::Storage::Cuda(c) => c, + _ => candle::bail!("hc_fused: {what} must be on CUDA"), + }; + // The (ptr, guard) pair borrows the slice, so bind before returning. + let ptr = match t.dtype() { + DType::F32 => { + let sl = cuda.as_cuda_slice::()?; + sl.slice(l.start_offset()..).device_ptr(sl.stream()).0 + } + DType::BF16 => { + let sl = cuda.as_cuda_slice::()?; + sl.slice(l.start_offset()..).device_ptr(sl.stream()).0 + } + other => candle::bail!("hc_fused: {what} has unsupported dtype {other:?}"), + }; + Ok(ptr as *const std::ffi::c_void) + } + + /// Allocate an output buffer of `dtype` and wrap it as a `Tensor`. + fn alloc_out( + dev: &candle_core::CudaDevice, + dtype: DType, + elems: usize, + shape: impl Into, + run: impl FnOnce(*mut std::ffi::c_void) -> i32, + ) -> Result { + use candle_core::cuda_backend::cudarc::driver::DevicePtr; + match dtype { + DType::F32 => { + let buf = unsafe { dev.alloc::(elems) }?; + let rc = run(buf.device_ptr(buf.stream()).0 as *mut std::ffi::c_void); + if rc != 0 { + candle::bail!("hc_fused: kernel dispatch refused (rc={rc})"); + } + let st = candle::CudaStorage::wrap_cuda_slice(buf, dev.clone()); + Ok(Tensor::from((candle::Storage::Cuda(st), shape.into()))) + } + DType::BF16 => { + let buf = unsafe { dev.alloc::(elems) }?; + let rc = run(buf.device_ptr(buf.stream()).0 as *mut std::ffi::c_void); + if rc != 0 { + candle::bail!("hc_fused: kernel dispatch refused (rc={rc})"); + } + let st = candle::CudaStorage::wrap_cuda_slice(buf, dev.clone()); + Ok(Tensor::from((candle::Storage::Cuda(st), shape.into()))) + } + other => candle::bail!("hc_fused: unsupported output dtype {other:?}"), + } + } + + /// Fused `hc_pre` tail: `y = sum_i pre[i] * x_f32[i, :]`, narrowed to + /// `out_dtype`. Replaces `broadcast_mul + sum(1) + to_dtype` (3 launches). + /// + /// - `x_f32` `[n, hc, h]` F32 contiguous (the promoted residual stack) + /// - `pre` `[n, hc]` F32 contiguous + pub fn hc_y_combine_cuda(x_f32: &Tensor, pre: &Tensor, out_dtype: DType) -> Result { + let (n, hc, h) = x_f32.dims3()?; + let (pn, phc) = pre.dims2()?; + if pn != n || phc != hc { + candle::bail!("hc_fused: pre is [{pn}, {phc}] but x_f32 is [{n}, {hc}, {h}]"); + } + if n == 0 || h == 0 { + candle::bail!("hc_fused: empty y-combine"); + } + let tag = dtype_tag(out_dtype)?; + let x_ptr = f32_ptr(x_f32, "x_f32")?; + let pre_ptr = f32_ptr(pre, "pre")?; + let dev = x_f32.device().as_cuda_device()?; + let stream = dev.cuda_stream().cu_stream() as i64; + let total = (n * h) as i64; + + #[allow(clippy::cast_possible_truncation)] + alloc_out(&dev, out_dtype, n * h, (n, h), |out| unsafe { + crate::cuda::ffi::hc_y_combine( + x_ptr, pre_ptr, out, hc as i32, h as i32, total, tag, stream, + ) + }) + } + + /// Fused `hc_post`: the whole 6-launch re-expansion in one launch. + /// + /// - `x` `[n, h]` model dtype + /// - `residual` `[n, hc, h]` model dtype + /// - `post` `[n, hc]` F32 + /// - `comb` `[n, hc, hc]` F32 + pub fn hc_post_fused_cuda( + x: &Tensor, + residual: &Tensor, + post: &Tensor, + comb: &Tensor, + ) -> Result { + let (n, h) = x.dims2()?; + let (rn, hc, rh) = residual.dims3()?; + if rn != n || rh != h { + candle::bail!("hc_fused: residual [{rn}, {hc}, {rh}] vs x [{n}, {h}]"); + } + if post.dims2()? != (n, hc) { + candle::bail!("hc_fused: post must be [{n}, {hc}], got {:?}", post.dims()); + } + if comb.dims3()? != (n, hc, hc) { + candle::bail!("hc_fused: comb must be [{n}, {hc}, {hc}], got {:?}", comb.dims()); + } + if x.dtype() != residual.dtype() { + candle::bail!( + "hc_fused: x is {:?} but residual is {:?}", + x.dtype(), + residual.dtype() + ); + } + if n == 0 || h == 0 { + candle::bail!("hc_fused: empty hc_post"); + } + let tag = dtype_tag(x.dtype())?; + let x_ptr = any_ptr(x, "x")?; + let res_ptr = any_ptr(residual, "residual")?; + let post_ptr = f32_ptr(post, "post")?; + let comb_ptr = f32_ptr(comb, "comb")?; + let dev = x.device().as_cuda_device()?; + let stream = dev.cuda_stream().cu_stream() as i64; + let total = (n * hc * h) as i64; + + #[allow(clippy::cast_possible_truncation)] + alloc_out(&dev, x.dtype(), n * hc * h, (n, hc, h), |out| unsafe { + crate::cuda::ffi::hc_post_fused( + x_ptr, res_ptr, post_ptr, comb_ptr, out, hc as i32, h as i32, total, tag, stream, + ) + }) + } } #[cfg(feature = "cuda")] -pub use cuda_impl::{hc_pre_fused_cuda, sqrt_softplus_cuda}; +pub use cuda_impl::{hc_post_fused_cuda, hc_pre_fused_cuda, hc_y_combine_cuda, sqrt_softplus_cuda}; + +#[cfg(not(feature = "cuda"))] +pub fn hc_y_combine_cuda( + _x_f32: &candle_core::Tensor, + _pre: &candle_core::Tensor, + _out_dtype: candle_core::DType, +) -> candle_core::Result { + candle_core::bail!("hc_y_combine_cuda requires the cuda feature") +} + +#[cfg(not(feature = "cuda"))] +pub fn hc_post_fused_cuda( + _x: &candle_core::Tensor, + _residual: &candle_core::Tensor, + _post: &candle_core::Tensor, + _comb: &candle_core::Tensor, +) -> candle_core::Result { + candle_core::bail!("hc_post_fused_cuda requires the cuda feature") +} // Non-CUDA stubs so call sites need no `cfg`. They are never reached: every // caller gates on [`usable`], which is false without a CUDA device. @@ -239,6 +420,139 @@ pub fn usable(t: &candle_core::Tensor) -> bool { && fused_enabled() } +/// Same gate as [`usable`] but for the model-dtype tensors (`hc_post`, the +/// `hc_pre` tail), which are BF16 rather than F32. +pub fn usable_model_dtype(t: &candle_core::Tensor) -> bool { + cfg!(feature = "cuda") + && t.device().is_cuda() + && matches!(t.dtype(), candle_core::DType::F32 | candle_core::DType::BF16) + && fused_enabled() +} + +// --------------------------------------------------------------------------- +// The on-GPU bitwise A/B — `ARC_HC_AB=1` +// --------------------------------------------------------------------------- +// +// hc_fused.cu's header has cited "the final proof is the on-GPU A/B: +// `ARC_HC_AB=1`" since the file was written. That A/B DID NOT EXIST — the +// string appeared nowhere in the tree except that sentence. This is it. +// +// It recomputes the eager candle chain beside every fused kernel and compares +// the two BIT FOR BIT (not by tolerance: -0.0 vs 0.0 and NaN payloads are +// differences, and a f32 subtraction would hide both). +// +// `ARC_HC_AB_POISON=1` perturbs the fused tensor by exactly one ULP before the +// comparison, so the check MUST go red. A guard that has never been observed +// failing is not evidence of anything — this codebase has shipped fifteen that +// passed on broken code, including the `#if defined(__USE_FAST_MATH__)` in +// hc_fused.cu, which nvcc can never define (measured; see that file's header). +use std::sync::atomic::{AtomicU64, Ordering}; + +static AB_CMPS: AtomicU64 = AtomicU64::new(0); +static AB_BAD_TENSORS: AtomicU64 = AtomicU64::new(0); +static AB_BAD_ELEMS: AtomicU64 = AtomicU64::new(0); + +pub fn ab_enabled() -> bool { + static ON: std::sync::OnceLock = std::sync::OnceLock::new(); + *ON.get_or_init(|| matches!(std::env::var("ARC_HC_AB").as_deref(), Ok("1"))) +} + +fn ab_poison() -> bool { + static ON: std::sync::OnceLock = std::sync::OnceLock::new(); + *ON.get_or_init(|| matches!(std::env::var("ARC_HC_AB_POISON").as_deref(), Ok("1"))) +} + +/// Raw bit patterns of a tensor, widened to u32 so BF16 and F32 share a path. +fn bits_of(t: &candle_core::Tensor) -> candle_core::Result> { + let t = t.flatten_all()?.contiguous()?; + Ok(match t.dtype() { + candle_core::DType::F32 => t.to_vec1::()?.iter().map(|v| v.to_bits()).collect(), + candle_core::DType::BF16 => t + .to_vec1::()? + .iter() + .map(|v| u32::from(v.to_bits())) + .collect(), + other => candle_core::bail!("hc_fused ab_check: unsupported dtype {other:?}"), + }) +} + +/// Compare a fused result against the eager one, bit for bit. +/// +/// Returns the number of mismatching elements. Logs the first mismatch loudly +/// and emits a periodic running total — the periodic line is deliberate: a +/// silent A/B is indistinguishable from an A/B that never ran, which is the +/// single most common way a green result in this repo has been wrong. +pub fn ab_check( + site: &str, + fused: &candle_core::Tensor, + eager: &candle_core::Tensor, +) -> candle_core::Result { + if fused.dims() != eager.dims() { + tracing::error!( + "HC_AB[{site}] SHAPE MISMATCH fused={:?} eager={:?}", + fused.dims(), + eager.dims() + ); + AB_BAD_TENSORS.fetch_add(1, Ordering::Relaxed); + return Ok(u64::MAX); + } + let mut fb = bits_of(fused)?; + let eb = bits_of(eager)?; + + if ab_poison() && !fb.is_empty() { + // Exactly one ULP up on the first element. If the comparison below + // still reports zero mismatches, the comparison is broken, not the + // kernel. + fb[0] = fb[0].wrapping_add(1); + } + + let mut bad = 0u64; + let mut first: Option<(usize, u32, u32)> = None; + for (i, (a, b)) in fb.iter().zip(eb.iter()).enumerate() { + if a != b { + bad += 1; + if first.is_none() { + first = Some((i, *a, *b)); + } + } + } + + let n = AB_CMPS.fetch_add(1, Ordering::Relaxed) + 1; + if bad > 0 { + AB_BAD_TENSORS.fetch_add(1, Ordering::Relaxed); + AB_BAD_ELEMS.fetch_add(bad, Ordering::Relaxed); + if let Some((i, a, b)) = first { + static ONCE: std::sync::Once = std::sync::Once::new(); + ONCE.call_once(|| { + tracing::error!( + "HC_AB[{site}] BITWISE MISMATCH: {bad}/{} elements differ; \ + first at [{i}] fused=0x{a:08x} eager=0x{b:08x}", + fb.len() + ); + }); + } + } + if n % 1000 == 0 { + tracing::info!( + "HC_AB: {n} comparisons, {} tensors with mismatches, {} bad elements \ + (poison={})", + AB_BAD_TENSORS.load(Ordering::Relaxed), + AB_BAD_ELEMS.load(Ordering::Relaxed), + ab_poison() + ); + } + Ok(bad) +} + +/// Running totals, for an end-of-run assertion. +pub fn ab_totals() -> (u64, u64, u64) { + ( + AB_CMPS.load(Ordering::Relaxed), + AB_BAD_TENSORS.load(Ordering::Relaxed), + AB_BAD_ELEMS.load(Ordering::Relaxed), + ) +} + /// Scalar f32 replicas that pin the fused kernels' op ORDER and ROUNDING to /// what the candle CUDA backend actually executes, bit for bit. /// @@ -428,6 +742,151 @@ pub(crate) mod reference { }) .collect() } + + /// The **candle op chain** of the `hc_pre` tail: + /// `pre.unsqueeze(-1).broadcast_mul(x_f32).sum(1).to_dtype(BF16)`. + /// + /// `x_f32` is `[hc][h]` row-major for one row `n`. Returns the BF16 bit + /// patterns, because the narrowing is part of what has to match. + /// F32 result, BEFORE the narrowing cast. The order assertions have to be + /// made here: BF16 keeps 8 mantissa bits, so narrowing discards almost + /// every reassociation difference and a bf16-only comparison would pass on + /// a kernel summing in the wrong order. (It did — the vacuity test below + /// caught exactly that and this split is the fix.) + pub(crate) fn hc_y_combine_candle_replay_f32( + x_f32: &[f32], + pre: &[f32], + hc: usize, + h: usize, + ) -> Vec { + (0..h) + .map(|col| { + // broadcast_mul materialises [hc, h] first... + let prods: Vec = (0..hc).map(|t| pre[t] * x_f32[t * h + col]).collect(); + // ...then `sum(1)` reduces the hc axis with block_dim = + // next_power_of_two(hc), which is candle_fast_sum's tree. + candle_fast_sum(&prods, super::candle_reduce_block_dim(hc)) + }) + .collect() + } + + /// The **fused kernel** `hc_y_combine_kernel`, transcribed, before narrowing. + pub(crate) fn hc_y_combine_fused_replay_f32( + x_f32: &[f32], + pre: &[f32], + hc: usize, + h: usize, + ) -> Vec { + (0..h) + .map(|col| { + let mut acc = vec![0.0f32; hc]; + for (t, a) in acc.iter_mut().enumerate() { + // the `shr[tid] = 0` start is not a no-op for -0.0f + *a = 0.0f32 + pre[t] * x_f32[t * h + col]; + } + let mut s = hc / 2; + while s > 0 { + for t in 0..s { + acc[t] += acc[t + s]; + } + s /= 2; + } + acc[0] + }) + .collect() + } + + /// Sequential `a0+a1+a2+a3` — what the kernel must NOT do. Exists only so + /// the vacuity test can show the comparison can tell the orders apart. + pub(crate) fn hc_y_combine_sequential_replay_f32( + x_f32: &[f32], + pre: &[f32], + hc: usize, + h: usize, + ) -> Vec { + (0..h) + .map(|col| { + let mut s = 0.0f32; + for t in 0..hc { + s += pre[t] * x_f32[t * h + col]; + } + s + }) + .collect() + } + + pub(crate) fn hc_y_combine_candle_replay( + x_f32: &[f32], + pre: &[f32], + hc: usize, + h: usize, + ) -> Vec { + hc_y_combine_candle_replay_f32(x_f32, pre, hc, h) + .into_iter() + .map(|s| half::bf16::from_f32(s).to_bits()) + .collect() + } + + /// The **fused kernel** `hc_y_combine_kernel`, transcribed. + pub(crate) fn hc_y_combine_fused_replay( + x_f32: &[f32], + pre: &[f32], + hc: usize, + h: usize, + ) -> Vec { + hc_y_combine_fused_replay_f32(x_f32, pre, hc, h) + .into_iter() + .map(|s| half::bf16::from_f32(s).to_bits()) + .collect() + } + + /// The **candle op chain** of `hc_post`, for one row `n`. + /// `x` is `[h]`, `residual` is `[hc][h]`, both already widened to f32. + pub(crate) fn hc_post_candle_replay( + x: &[f32], + residual: &[f32], + post: &[f32], + comb: &[f32], + hc: usize, + h: usize, + ) -> Vec { + let mut out = Vec::with_capacity(hc * h); + for k in 0..hc { + for (col, &xv) in x.iter().enumerate().take(h) { + let term1 = post[k] * xv; + // term2 = sum_j comb[j][k] * residual[j][col] — the K=hc GEMM. + let mut term2 = 0.0f32; + for j in 0..hc { + term2 = comb[j * hc + k].mul_add(residual[j * h + col], term2); + } + out.push(half::bf16::from_f32(term1 + term2).to_bits()); + } + } + out + } + + /// The **fused kernel** `hc_post_fused_kernel`, transcribed. + pub(crate) fn hc_post_fused_replay( + x: &[f32], + residual: &[f32], + post: &[f32], + comb: &[f32], + hc: usize, + h: usize, + ) -> Vec { + let mut out = vec![0u16; hc * h]; + for k in 0..hc { + for col in 0..h { + let t1 = post[k] * x[col]; + let mut acc = 0.0f32; + for j in 0..hc { + acc = comb[j * hc + k].mul_add(residual[j * h + col], acc); + } + out[k * h + col] = half::bf16::from_f32(t1 + acc).to_bits(); + } + } + out + } } #[cfg(test)] @@ -444,19 +903,35 @@ mod tests { } } - /// Source-level tripwires, mirroring `sinkhorn::tests`: the IEEE - /// intrinsics and the fast-math `#error` guard must stay in hc_fused.cu, - /// fast-math approximations must stay out, and build.rs must keep the file - /// out of the `--use_fast_math` glob and in the `--fmad=false` builder. - /// Getting this wiring wrong is silent — the kernel still runs, it just - /// stops being bit-identical — so the `#error` guard is the hard stop and - /// these string checks catch it on CPU CI too. + /// Source-level tripwires, mirroring `sinkhorn::tests`: the IEEE intrinsics + /// must stay in hc_fused.cu, fast-math approximations must stay out, and + /// build.rs must keep the file out of the `--use_fast_math` glob and in the + /// `--fmad=false` builder. Getting this wiring wrong is silent — the kernel + /// still runs, it just stops being bit-identical. + /// + /// NOTE ON WHAT ACTUALLY PROTECTS THIS FILE. The `#if + /// defined(__USE_FAST_MATH__)` `#error` in hc_fused.cu (and in sinkhorn.cu) + /// is DECORATIVE: nvcc 12.4 defines neither `__USE_FAST_MATH__` nor + /// `__FAST_MATH__` in either the host or the device pass, so it cannot fire. + /// Measured, not assumed — + /// nvcc --use_fast_math -E -dM x.cu | grep -i fast -> no output + /// nvcc --use_fast_math -arch=sm_90 -c guard.cu -> compiles clean + /// while `--use_fast_math` does reach the device pass (a float divide drops + /// from 3 MUFU/RCP instructions to 1). The real protection is the build.rs + /// wiring asserted below, plus the runtime `ARC_HC_AB=1` bitwise A/B, which + /// catches a mis-compiled build on the first token. The presence assertion + /// is kept only so the comment above the guard travels with it. #[test] fn kernel_source_and_build_wiring_guards() { let cu = include_str!("hc_fused.cu"); assert!( cu.contains("#if defined(__USE_FAST_MATH__)") && cu.contains("#error"), - "hc_fused.cu lost its fast-math #error guard" + "hc_fused.cu lost its (decorative) fast-math #error guard" + ); + assert!( + cu.contains("THE COMPILE-TIME FAST-MATH GUARD BELOW IS DEAD"), + "hc_fused.cu lost the note recording that the #error guard cannot fire; \ + without it the next reader will trust a guard that has never been able to fail" ); for required in ["__fadd_rn", "__fmul_rn", "candle_recip", "candle_sigmoid"] { assert!(cu.contains(required), "hc_fused.cu lost required token {required}"); @@ -635,4 +1110,99 @@ mod tests { ); } } + + /// Inputs shaped like V4 decode: hc_mult = 4, one row, a slice of hidden. + fn y_combine_inputs(h: usize) -> (Vec, Vec) { + let hc = 4; + let mut r = Lcg(0x5eed_1234); + let x: Vec = (0..hc * h).map(|_| r.next_f32(-4.0, 4.0)).collect(); + let pre: Vec = (0..hc).map(|_| r.next_f32(0.0, 1.0)).collect(); + (x, pre) + } + + #[test] + fn hc_y_combine_fused_is_bit_identical_to_candle_chain() { + let (hc, h) = (4usize, 257usize); + let (x, pre) = y_combine_inputs(h); + // Assert at F32, before the narrowing: BF16 keeps 8 mantissa bits and + // would hide an order difference (see the vacuity test below). + assert_eq!( + hc_y_combine_candle_replay_f32(&x, &pre, hc, h), + hc_y_combine_fused_replay_f32(&x, &pre, hc, h), + "fused y-combine diverged from the candle sum(1) chain in F32" + ); + // ...and again after it, which is what the kernel actually stores. + assert_eq!( + hc_y_combine_candle_replay(&x, &pre, hc, h), + hc_y_combine_fused_replay(&x, &pre, hc, h), + "fused y-combine diverged after the BF16 narrowing" + ); + } + + /// The pairwise tree is the whole point: candle's `sum(1)` over hc = 4 gives + /// `(a0+a2)+(a1+a3)`, not `((a0+a1)+a2)+a3`. + /// + /// This test previously compared the two orders AFTER narrowing to BF16 and + /// FAILED — not because the kernel was wrong, but because 8 mantissa bits + /// swallow essentially every reassociation difference, which meant the + /// bit-identity test above would have passed on a kernel summing in the + /// wrong order. That is the failure this vacuity check exists to expose, and + /// the fix is to make the claim where it is observable: at F32. + #[test] + fn hc_y_combine_tree_order_is_not_vacuous() { + let (hc, h) = (4usize, 257usize); + let (x, pre) = y_combine_inputs(h); + let tree = hc_y_combine_candle_replay_f32(&x, &pre, hc, h); + let sequential = hc_y_combine_sequential_replay_f32(&x, &pre, hc, h); + let differing = tree + .iter() + .zip(sequential.iter()) + .filter(|(a, b)| a.to_bits() != b.to_bits()) + .count(); + assert!( + differing > 0, + "sequential summation produced bit-identical F32 output on all {h} elements, so \ + the tree-order assertion proves nothing on this data" + ); + } + + fn post_inputs(h: usize) -> (Vec, Vec, Vec, Vec) { + let hc = 4; + let mut r = Lcg(0xabcd_ef01); + let x: Vec = (0..h).map(|_| r.next_f32(-3.0, 3.0)).collect(); + let residual: Vec = (0..hc * h).map(|_| r.next_f32(-3.0, 3.0)).collect(); + let post: Vec = (0..hc).map(|_| r.next_f32(0.0, 2.0)).collect(); + let comb: Vec = (0..hc * hc).map(|_| r.next_f32(0.0, 1.0)).collect(); + (x, residual, post, comb) + } + + #[test] + fn hc_post_fused_is_bit_identical_to_candle_chain() { + let (hc, h) = (4usize, 193usize); + let (x, residual, post, comb) = post_inputs(h); + let a = hc_post_candle_replay(&x, &residual, &post, &comb, hc, h); + let b = hc_post_fused_replay(&x, &residual, &post, &comb, hc, h); + assert_eq!(a, b, "fused hc_post diverged from the candle term1+term2 chain"); + } + + /// `term2` transposes `comb` before the matmul. Reading `comb[k][j]` instead + /// of `comb[j][k]` is the single easiest way to get this kernel wrong while + /// keeping every shape valid, so prove the comparison would catch it. + #[test] + fn hc_post_guard_catches_a_transposed_comb() { + let (hc, h) = (4usize, 193usize); + let (x, residual, post, comb) = post_inputs(h); + let reference = hc_post_candle_replay(&x, &residual, &post, &comb, hc, h); + let mut transposed = vec![0.0f32; hc * hc]; + for j in 0..hc { + for k in 0..hc { + transposed[j * hc + k] = comb[k * hc + j]; + } + } + let mutated = hc_post_fused_replay(&x, &residual, &post, &transposed, hc, h); + assert!( + reference.iter().zip(mutated.iter()).any(|(a, b)| a != b), + "a transposed comb produced bit-identical output, so the hc_post assertion is vacuous" + ); + } } diff --git a/mistralrs-core/src/models/dsv4_mhc.rs b/mistralrs-core/src/models/dsv4_mhc.rs index 57fb1e2d7..f7f8d6513 100644 --- a/mistralrs-core/src/models/dsv4_mhc.rs +++ b/mistralrs-core/src/models/dsv4_mhc.rs @@ -392,8 +392,36 @@ impl V4MHCLayerParams { // a contiguous tensor is metadata-only: no kernel, no allocation. The // values are bit-identical — BF16->F32 is an exact widening. let x_f32 = x_flat.reshape((n, hc, h))?; // [N, hc, h], F32 - let pre_b = pre.unsqueeze(D::Minus1)?; // [N, hc, 1] - let y = pre_b.broadcast_mul(&x_f32)?.sum(1)?; // [N, hidden] + + // The eager spelling of the tail is THREE launches — a broadcast_mul + // over [N, hc, hidden], a `sum(1)` over the 4-element hc axis, and the + // narrowing cast — to compute a 4-term weighted average. Fused into one + // kernel that replays candle's `fast_sum` pairwise tree exactly; see + // `hc_y_combine_kernel` in cuda/hc_fused.cu for the order it replays and + // why it is not the sequential sum. + let eager_y = || -> Result { + let pre_b = pre.unsqueeze(D::Minus1)?; // [N, hc, 1] + pre_b + .broadcast_mul(&x_f32)? + .sum(1)? // [N, hidden] F32 + .to_dtype(in_dtype) + }; + // `fused` alone is not enough: the tail also has to NARROW to the model + // dtype, and only F32/BF16 rounding has been transcribed. An unlisted + // dtype takes the eager path rather than erroring or, worse, rounding + // differently. + let tail_fused = fused + && crate::cuda::hc_fused::fuse2_enabled() + && matches!(in_dtype, DType::F32 | DType::BF16); + let y_nd = if tail_fused { + let out = crate::cuda::hc_fused::hc_y_combine_cuda(&x_f32, &pre, in_dtype)?; + if crate::cuda::hc_fused::ab_enabled() { + crate::cuda::hc_fused::ab_check("hc_pre.y", &out, &eager_y()?)?; + } + out + } else { + eager_y()? + }; // Restore leading dims. let mut y_shape = leading_dims.clone(); @@ -404,9 +432,10 @@ impl V4MHCLayerParams { comb_shape.push(hc); comb_shape.push(hc); - // `y` feeds layernorm + attention/MLP, so it goes back to the model - // dtype. `post` / `comb` stay F32 — see the dtype contract above. - let y_out = y.reshape(y_shape)?.to_dtype(in_dtype)?; + // `y` is already in the model dtype (both paths narrow it), so this + // reshape is metadata-only. `post` / `comb` stay F32 — see the dtype + // contract above. + let y_out = y_nd.reshape(y_shape)?; let post_out = post.reshape(post_shape)?; let comb_out = comb.reshape(comb_shape)?; Ok((y_out, post_out, comb_out)) @@ -426,6 +455,13 @@ impl V4MHCLayerParams { /// below are retained only for hand-constructed callers and short-circuit /// to a clone (no kernel) on the production path. /// + /// On CUDA the whole body is ONE kernel (`hc_post_fused`), which is + /// bit-identical to the eager chain kept below it — verified on live decode + /// data, not argued: `ARC_HC_AB=1` compared 4,000 tensors with zero + /// mismatching bits, and the same comparison flags all 4,000 when + /// `ARC_HC_AB_POISON=1` perturbs one element by a single ULP. Opt out with + /// `ARC_HC_FUSE2=0`. + /// /// Returns the new residual stack with the same shape *and dtype* as `x` /// (the branch output), i.e. the model dtype. pub fn hc_post( @@ -449,31 +485,58 @@ impl V4MHCLayerParams { ), }; - let x_n = x.reshape((n, h))?.to_dtype(DType::F32)?; - let residual_n = residual - .reshape((n, hc, h))? - .to_dtype(DType::F32)?; - let post_n = post.reshape((n, hc))?.to_dtype(DType::F32)?; - let comb_n = comb.reshape((n, hc, hc))?.to_dtype(DType::F32)?; - - // term1 = post[..., :, None] * x[..., None, :] → [N, hc, h] - let term1 = post_n - .unsqueeze(D::Minus1)? // [N, hc, 1] - .broadcast_mul(&x_n.unsqueeze(1)?)?; // [N, hc, h] + let x_2d = x.reshape((n, h))?.contiguous()?; + let residual_3d = residual.reshape((n, hc, h))?.contiguous()?; + let post_2d = post.reshape((n, hc))?.to_dtype(DType::F32)?.contiguous()?; + let comb_3d = comb + .reshape((n, hc, hc))? + .to_dtype(DType::F32)? + .contiguous()?; - // term2 = sum_j comb[..., j, k] * residual[..., j, :] - // = einsum("njk,njh->nkh", comb, residual) - // Implementation: for each k, term2[n, k, h] = sum_j comb[n, j, k] * residual[n, j, h] - // Use matmul: comb.transpose(1,2) @ residual → [N, hc, h] - let term2 = comb_n.transpose(1, 2)?.matmul(&residual_n)?; // [N, hc, h] + // The eager spelling is SIX launches — two promoting casts, a + // broadcast_mul, a transpose that candle services with a real copy, the + // [hc,hc]x[hc,h] matmul, the add, and the narrowing cast — for four + // multiply-adds per output element. At b=1 that is 516 of the decode + // step's launches across the 86 mixing points. + // + // The fused kernel is bit-identical by construction: BF16->F32 widening + // is exact, the K=hc accumulation is the same sequential FMA a K=4 GEMM + // performs, and the narrowing uses candle's own `__float2bfloat16`. + // `ARC_HC_AB=1` checks that claim against this very chain on live data. + let eager = || -> Result { + let x_n = x_2d.to_dtype(DType::F32)?; + let residual_n = residual_3d.to_dtype(DType::F32)?; + // term1 = post[..., :, None] * x[..., None, :] → [N, hc, h] + let term1 = post_2d + .unsqueeze(D::Minus1)? // [N, hc, 1] + .broadcast_mul(&x_n.unsqueeze(1)?)?; // [N, hc, h] + // term2 = sum_j comb[..., j, k] * residual[..., j, :] + // = einsum("njk,njh->nkh", comb, residual) + let term2 = comb_3d.transpose(1, 2)?.matmul(&residual_n)?; // [N, hc, h] + (term1 + term2)?.to_dtype(in_dtype) + }; - let out_n = (term1 + term2)?; + let fused = crate::cuda::hc_fused::fuse2_enabled() + && crate::cuda::hc_fused::usable_model_dtype(&x_2d) + && x_2d.dtype() == residual_3d.dtype() + && crate::cuda::hc_fused::usable(&post_2d); + let out_n = if fused { + let out = + crate::cuda::hc_fused::hc_post_fused_cuda(&x_2d, &residual_3d, &post_2d, &comb_3d)?; + if crate::cuda::hc_fused::ab_enabled() { + crate::cuda::hc_fused::ab_check("hc_post", &out, &eager()?)?; + } + out + } else { + eager()? + }; - // Restore shape. + // Restore shape. `out_n` is already in the model dtype on both paths, so + // this is metadata-only. let mut out_shape = leading_dims; out_shape.push(hc); out_shape.push(h); - out_n.reshape(out_shape)?.to_dtype(in_dtype) + out_n.reshape(out_shape) } /// Spec convenience wrapper: full attn-side round-trip for one decoder layer. From cd647cedb9c317cbbec69ab5b3439a4d549f31b4 Mon Sep 17 00:00:00 2001 From: Arc Agent Date: Thu, 20 Aug 2026 05:18:22 +0000 Subject: [PATCH 24/27] perf(ArcInfer/ArcAttention): fuse the V4 Q/K pre-attention block -- 16 launches per layer become 1 The b=1 decode step is launch-bound, and after the mHC/router fusions 34.3% of what is left is PURE DATA MOVEMENT: ucopy_bf16 514/step, copy2d_bf16 418, cast_bf16_f32 356, cast_f32_bf16 221, cast_u32_f32 203. This takes the largest ucopy/copy2d contributor in that list. `Attention::forward` spelled the head transpose, the per-head Q RMS-norm and the RoPE as SIXTEEN candle launches per layer, TEN of them pure copies materialising intermediates that exist only because the chain is written as separate ops: a transpose `.contiguous()`, two `.contiguous()` narrows inside `rope_i`, and two two-input `cat`s. Their entire data footprint at decode is 64 KB of Q and 1 KB of K. `cuda/qk_norm_rope.cu` does all of it in one launch of n_heads+1 blocks. The transpose becomes an output ADDRESS, the NoPE/PE split and re-`cat` become two write ranges of the same row, and the RMS statistic stays in a register. The kv projection is hoisted above the Q RMS-norm so both rows fit one launch; the two projections read only `xs`, so no arithmetic moves. BIT-IDENTICAL, proven on the live model rather than asserted: ARC_QK_VERIFY=1 runs both paths at every layer of every step and compares bf16 BIT PATTERNS, and every comparison first runs a negative control -- a 1-ULP poison that MUST be detected, else the leg declares itself VOID rather than passed. 16 forwards x 43 layers x 2 tensors, zero mismatches, on an H200 with clocks locked at 1980 MHz. Getting there cost two wrong theories, both recorded in the kernel header so the next reader does not re-derive them: * bf16 arithmetic is NOT immune to --fmad. nvcc contracts one of the two products of `a*c - b*s` into an fma.rn.bf16, rounding it once with the subtraction instead of twice. Measured for head 0 / pair 0 at position 1 (a=-5.65625, b=-3.546875, cos=0x3f0a, sin=0x3f57, with the diagnostic confirming both sides read the SAME table row): three roundings give -0.0625, full f32 gives -0.069824, and candle gives -0.064453125 -- an exact tie resolved by round-half-to-even. So the file gets its OWN builder carrying candle-kernels' flags, and the contraction is written out with __hfma rather than left to the compiler. * The first attempt at that fix changed NOTHING -- byte-identical output. `ar t` showed the stale --fmad=false object still inside libmistralrssinkhornieee.a while the new archive was empty, so the linker resolved the old definition and a green build proved nothing. The symbol is renamed arc_qk_norm_rope_bf16_v2 so that failure mode becomes a link error, never a silent wrong answer. Toggle with ARC_QK_FUSED=0 (restores the eager chain, so the A/B runs from ONE binary and cannot be confounded by a rebuild). The shape gate refuses anything outside the specialised set rather than approximating it, and logs the refusal once, so a permanently disengaged fast path cannot masquerade as a working one. An engagement counter backs that up: "the benchmark got faster" and "the kernel ran" are different claims. Co-Authored-By: Claude Opus 5 (1M context) MEASURED, H200, ONE binary toggled with ARC_QK_FUSED, clocks stable at 1830 MHz in every leg (that is this card's actual load clock; `-lgc 1980` is applied but nothing throttles -- 122 W of 700 W, 28 C, no active clock-event reasons). ops/token 5,347.8 -> 4,744.9 launches/forward (-602.9, -11.3%) memcpy 1,840.3 -> 1,580.9 (-259.4, -14.1%) allocs 212.7 -> 205.7 ( -7.1, -3.3%) Both legs pinned by TWO INDEPENDENT ANCHORS that had to agree before the leg counted: the API's completion_tokens, and a per-layer kernel count factoring as 43 x k x N (bminimum_f32 = 27,520 = 43 x 4 x 160). Counts read from a full GROUP BY over the exported sqlite, never a top-N table. Every kernel that moved is one of the sixteen this change removes, and the per-kernel deltas reconcile to the launch: ucopy_bf16 474.1 -> 300.8 -173.3 (~4 x 43) rope_i_bf16 155.5 -> 69.5 -86.0 ( 2 x 43) affine_bf16 128.7 -> 42.7 -86.0 ( 2 x 43) copy2d_bf16 376.4 -> 290.9 -85.5 (~2 x 43) bmul_bf16 43.0 -> 0.0 -43.0 fast_sum_bf16 43.0 -> 0.0 -43.0 urecip_bf16 43.0 -> 0.0 -43.0 usqr_bf16 43.0 -> 0.0 -43.0 usqrt_bf16 43.0 -> 0.0 -43.0 qk_norm_rope_kernel 0.0 -> 43.0 +43.0 ------- -602.8 vs -602.9 measured Decode, INTERLEAVED A-B-A-B (a single before/after on a shared box is not evidence), ms/token by the difference method so TTFT cancels: OFF r1 37.967 ms 26.34 tok/s ON r1 35.732 ms 27.99 tok/s OFF r2 37.960 ms 26.34 tok/s ON r2 34.964 ms 28.60 tok/s ---------------------------------------------------------------- OFF 37.964 ms 26.34 tok/s ON 35.348 ms 28.29 tok/s -6.9% ms/token, +7.4% tok/s. The two OFF legs agree to 0.02%, and neither ON leg overlaps either OFF leg. The OFF baseline reproduces the previous session's 38.18 ms/token independently. All four legs returned the SAME 160 greedy tokens (text sha 6d6456512d885a6f), so 640 tokens are byte-identical across the toggle -- an end-to-end check on top of the per-tensor bit comparison. Engagement asserted in both directions: engaged_lines=1 / declined_lines=0 with the kernel on, and 0/0 with it off. Co-Authored-By: Claude Opus 5 (1M context) --- mistralrs-core/build.rs | 41 ++- mistralrs-core/src/cuda/ffi.rs | 27 ++ mistralrs-core/src/cuda/mod.rs | 1 + mistralrs-core/src/cuda/qk_norm_rope.cu | 355 +++++++++++++++++++ mistralrs-core/src/cuda/qk_norm_rope.rs | 445 ++++++++++++++++++++++++ mistralrs-core/src/layers.rs | 13 + mistralrs-core/src/models/deepseek4.rs | 243 +++++++++++-- 7 files changed, 1092 insertions(+), 33 deletions(-) create mode 100644 mistralrs-core/src/cuda/qk_norm_rope.cu create mode 100644 mistralrs-core/src/cuda/qk_norm_rope.rs diff --git a/mistralrs-core/build.rs b/mistralrs-core/build.rs index bfe9f4226..e2243c12e 100644 --- a/mistralrs-core/build.rs +++ b/mistralrs-core/build.rs @@ -10,7 +10,7 @@ fn main() { println!("cargo:rerun-if-changed=build.rs"); let build_dir = PathBuf::from(std::env::var("OUT_DIR").unwrap()); - // sinkhorn.cu and hc_fused.cu are EXCLUDED from this fast-math builder + // sinkhorn.cu, hc_fused.cu and qk_norm_rope.cu 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 -> @@ -18,7 +18,7 @@ fn main() { // being re-globbed under fast math. See mistralrs-core/src/cuda/. let mut builder = cudaforge::KernelBuilder::new() .source_glob("src/cuda/*.cu") - .exclude(&["sinkhorn.cu", "hc_fused.cu"]) + .exclude(&["sinkhorn.cu", "hc_fused.cu", "qk_norm_rope.cu"]) .out_dir(&build_dir) .arg("-std=c++17") .arg("-O3") @@ -96,6 +96,43 @@ fn main() { println!("cargo:rustc-link-search={}", sinkhorn_dir.display()); println!("cargo:rustc-link-lib=mistralrssinkhornieee"); + // Third builder: qk_norm_rope.cu must be bit-identical to candle's + // `ropei`, and candle-kernels compiles with `-O3 + // --expt-relaxed-constexpr -std=c++17` and the DEFAULT `-fmad=true`. + // Under those flags nvcc contracts one of the two products of + // `a*c - b*s` into an `fma.rn.bf16`, which rounds once instead of + // twice. Measured on an H200: with `--fmad=false` that element comes + // out -0.0625 and candle gives -0.064453125 (an exact tie resolved by + // round-half-to-even), so 306 of 4,096 RoPE outputs per layer differ. + // It therefore CANNOT join the IEEE builder above -- sinkhorn.cu and + // hc_fused.cu need `--fmad=false` to forbid contraction; this file + // needs contraction, because that is what it is reproducing. It stays + // out of the fast-math glob for the usual reason. + let qk_dir = build_dir.join("qk_candlematch"); + let mut qk_builder = cudaforge::KernelBuilder::new() + .source_files(vec!["src/cuda/qk_norm_rope.cu"]) + .out_dir(&qk_dir) + .arg("-std=c++17") + .arg("-O3") + .arg("--expt-relaxed-constexpr") + .arg("--verbose") + .arg("--compiler-options") + .arg("-fPIC"); + if let Some(cuda_nvcc_flags_env) = CUDA_NVCC_FLAGS { + qk_builder = qk_builder.arg("--compiler-options"); + qk_builder = qk_builder.arg(cuda_nvcc_flags_env); + } + let qk_out = if target.contains("msvc") { + qk_dir.join("mistralrsqkcandlematch.lib") + } else { + qk_dir.join("libmistralrsqkcandlematch.a") + }; + qk_builder + .build_lib(qk_out) + .expect("Build qk_norm_rope kernel failed!"); + println!("cargo:rustc-link-search={}", qk_dir.display()); + println!("cargo:rustc-link-lib=mistralrsqkcandlematch"); + println!("cargo:rustc-link-lib=dylib=cudart"); if target.contains("msvc") { diff --git a/mistralrs-core/src/cuda/ffi.rs b/mistralrs-core/src/cuda/ffi.rs index 441e0d954..ade2d091b 100644 --- a/mistralrs-core/src/cuda/ffi.rs +++ b/mistralrs-core/src/cuda/ffi.rs @@ -375,6 +375,33 @@ extern "C" { // are compiled by build.rs's dedicated no-fast-math builder; see the // bit-identity contract at the top of that file. #[allow(clippy::too_many_arguments)] + /// Fused DeepSeek-V4 Q/K pre-attention block (`cuda/qk_norm_rope.cu`): + /// head transpose + per-head Q RMS-norm + adjacent-pair RoPE + NoPE/PE + /// recombination, in one launch. Returns 0 on launch, non-zero if the + /// (dtype, head_dim, rope_dim) triple is outside the specialised set — the + /// caller must then run the eager chain rather than approximate. + #[allow(clippy::too_many_arguments)] + pub(crate) fn arc_qk_norm_rope_bf16_v2( + q_in: *const c_void, + k_in: *const c_void, + cos_tab: *const c_void, + sin_tab: *const c_void, + q_out: *mut c_void, + k_out: *mut c_void, + n_heads: i32, + batch: i32, + seq_len: i32, + head_dim: i32, + rope_dim: i32, + pos_offset: i32, + dtype: i32, + inv_n_bits: u32, + zero_bits: u32, + one_bits: u32, + eps_bits: u32, + stream: i64, + ) -> i32; + pub(crate) fn hc_pre_fused_f32( x_flat: *const c_void, mixes_raw: *const c_void, diff --git a/mistralrs-core/src/cuda/mod.rs b/mistralrs-core/src/cuda/mod.rs index 68f0e65ff..c60f89aa5 100644 --- a/mistralrs-core/src/cuda/mod.rs +++ b/mistralrs-core/src/cuda/mod.rs @@ -2,5 +2,6 @@ pub mod ffi; pub mod gdn; pub mod hc_fused; pub mod moe; +pub mod qk_norm_rope; pub mod sinkhorn; pub mod ssm; diff --git a/mistralrs-core/src/cuda/qk_norm_rope.cu b/mistralrs-core/src/cuda/qk_norm_rope.cu new file mode 100644 index 000000000..65f0c9bbf --- /dev/null +++ b/mistralrs-core/src/cuda/qk_norm_rope.cu @@ -0,0 +1,355 @@ +// Parent system: ArcInfer / ArcAttention +// +// Fused DeepSeek-V4 Q/K pre-attention block: head-transpose + per-head Q +// RMS-normalisation + adjacent-pair RoPE + the NoPE/PE recombination, in ONE +// kernel. +// +// --------------------------------------------------------------------------- +// WHY +// --------------------------------------------------------------------------- +// At b=1 the V4 decode step is launch-bound. After the mHC/router fusions the +// step is 5,566 launches, of which 1,910.8 (34.3%) are PURE DATA MOVEMENT: +// ucopy_bf16 514, copy2d_bf16 418, cast_bf16_f32 356, cast_f32_bf16 221, +// cast_u32_f32 203. They move kilobytes and cost microseconds each. +// +// The Q/K pre-attention block in `deepseek4.rs::Attention::forward` is the +// largest single contributor of `ucopy`/`copy2d` in that list. It spells a +// transpose, an RMS-norm and a RoPE as SIXTEEN candle launches per layer, of +// which TEN are pure copies: +// +// q.reshape.transpose(1,2).contiguous() 1 copy (materialised transpose) +// q.sqr 1 +// .mean_keepdim = fast_sum + affine(1/D) 2 +// .affine(1.0, eps) 1 +// .recip 1 +// .sqrt 1 +// q.broadcast_mul(inv_rms) 1 +// k.reshape.transpose(1,2).contiguous() 1 copy (free at T=1) +// q_pe.contiguous() (inside rope_i) 1 copy +// rope_i(q_pe) 1 +// k_pe.contiguous() (inside rope_i) 1 copy +// rope_i(k_pe) 1 +// cat([q_nope, q_pe]) 2 copies +// cat([k_nope, k_pe]) 2 copies +// +// Sixteen launches whose entire data footprint at decode is 64 KB of Q and +// 1 KB of K. This file does the whole thing in one launch of 65 blocks. +// +// The intermediate tensors that the copies exist to materialise are never +// built: the transpose is expressed as an output ADDRESS, the NoPE/PE split +// and re-`cat` as two write ranges of the same row, and the RMS statistic is +// kept in a register between the reduction and the scaling. +// +// --------------------------------------------------------------------------- +// BIT-IDENTITY CONTRACT +// --------------------------------------------------------------------------- +// Q feeds attention scores and K is written into the KV cache, so a +// reassociated sum or a contracted FMA does not merely perturb a logit — it +// perturbs every future token that reads that cache entry. This kernel is +// therefore BIT-IDENTICAL to the candle chain, not an approximation of it, and +// `ARC_QK_VERIFY=1` proves that on real tensors at every layer of every step, +// with a negative control. +// +// Three things make that achievable rather than merely hoped for: +// +// 1. **The reduction tree shape is copied, not chosen.** candle's `fast_sum` +// (candle-kernels/src/reduce.cu) accumulates in T -- NOT in f32; see the +// standing TODO at the top of that file -- into `__shared__ T shr[]`, one +// element per thread, then halves with `for (s = blockDim.x/2; s > 0; +// s >>= 1)`. The result depends on `blockDim.x`, which candle sets to +// `min(1024, el_to_sum_per_block).next_power_of_two()`. `HEAD_DIM` here is +// that same block width (asserted host-side in `qk_norm_rope.rs`), so the +// tree has the same shape and the same rounding. A row length whose block +// width does not equal the row length is REFUSED, not approximated. +// +// 2. **The FMA contraction is REPRODUCED, not avoided.** This was got wrong +// once and the `ARC_QK_VERIFY` harness caught it, so the correction is +// recorded here with its measurement. bf16 arithmetic is NOT immune to +// `--fmad`: nvcc contracts one of the two products of `a*c - b*s` into an +// `fma.rn.bf16`, so that product is rounded once with the subtraction +// instead of twice. Measured on an H200 for head 0 / pair 0 at position 1 +// (a = -5.65625, b = -3.546875, cos = 0x3f0a, sin = 0x3f57, and the +// diagnostic confirmed BOTH sides read the same table row): +// +// round(round(a*c) - round(b*s)) = -0.062500 (--fmad=false) +// round_f32(a*c - b*s) = -0.069824 (full f32) +// round(a*c - round(b*s)) = -0.064453125 (CANDLE) +// +// The third is an exact tie -- significand 1.02734375, midway between +// 1.0234375 and 1.03125 -- resolved by round-half-to-even. Building this +// file with `--fmad=false` therefore made 306 of the 4,096 RoPE outputs +// per layer differ, all on the low-j (high-frequency) pairs, because +// high-j pairs have cos ~ 1 / sin ~ 0 and round the same either way. +// `build.rs` consequently gives this file its OWN builder carrying +// candle-kernels' exact flags (`-O3 --expt-relaxed-constexpr -std=c++17`, +// default `-fmad=true`, no fast math) rather than the IEEE builder that +// sinkhorn.cu and hc_fused.cu need. The f32 path is still REFUSED below: +// the same reasoning would have to be re-established for f32's own +// contraction behaviour, and it has not been measured. +// +// 3. **The scalars are converted host-side by the same code candle uses.** +// candle's `affine(mul, add)` narrows f64 -> T before the kernel sees it; +// `1/512` and `rms_norm_eps` arrive here as bf16 bit patterns produced by +// `half::bf16::from_f64`, so no rounding decision is taken twice. +// +// This file is compiled by the DEDICATED IEEE BUILDER in build.rs, alongside +// sinkhorn.cu and hc_fused.cu, and carries the same guard against being +// re-globbed under `--use_fast_math`. +// --------------------------------------------------------------------------- + +#ifdef __CUDA_ARCH__ +#if defined(__CUDACC_RTC__) +#error "qk_norm_rope.cu must not be compiled with NVRTC" +#endif +#endif + +// Guard against this file being moved back under the fast-math glob in +// build.rs. `--use_fast_math` defines __CUDA_FAST_MATH__ and would rewrite the +// f32 helpers used by hsqrt/__hdiv's approximation sandwich. +#ifdef __CUDA_FAST_MATH__ +#error "qk_norm_rope.cu is bit-identity critical and must be built WITHOUT --use_fast_math (see build.rs: it belongs to the IEEE builder, not the fast-math glob)" +#endif + +#include +#include +#include +#include +#include +#include + +// dtype tags shared with qk_norm_rope.rs. Only BF16 is accepted; see (2) above. +#define QK_DTYPE_BF16 1 + +namespace { + +// candle `urecip_bf16` == `recipg(x)` == `__nv_bfloat16 one = 1.0; one / a` +// (candle-kernels/src/cuda_utils.cuh). `operator/` is `__hdiv`, a div.rn.f32 +// sandwich. Spelled the same way here so the same PTX is emitted. +__device__ __forceinline__ __nv_bfloat16 qk_recip(__nv_bfloat16 a) { + __nv_bfloat16 one = 1.0; + return one / a; +} + +// candle `usqrt_bf16` == `sqrtg(x)` == `hsqrt(a)`. +__device__ __forceinline__ __nv_bfloat16 qk_sqrt(__nv_bfloat16 a) { return hsqrt(a); } + +} // namespace + +// --------------------------------------------------------------------------- +// The fused kernel. +// +// Grid : (H_Q + 1, B * T) -- one block per output row; the +1 row is K. +// Block : HEAD_DIM threads -- one thread per element, which is also exactly +// candle's fast_sum block width for this row +// length, so the reduction tree matches. +// +// Q in : [B, T, H, D] contiguous, straight off the q_proj GEMM (NOT +// transposed -- the transpose is this kernel's output addressing). +// K in : [B, T, 1, D] contiguous, straight off kv_norm. +// Q out : [B, H, T, D] contiguous. +// K out : [B, 1, T, D] contiguous. +// +// `cos`/`sin` are the FULL [max_pos, D_rope/2] tables; `pos_offset` is +// `seqlen_offsets[0]`, i.e. candle's `self.cos.narrow(0, offset, seq_len)` +// expressed as a base offset instead of a view. +// --------------------------------------------------------------------------- +template +__global__ void qk_norm_rope_kernel( + const __nv_bfloat16 *__restrict__ q_in, + const __nv_bfloat16 *__restrict__ k_in, + const __nv_bfloat16 *__restrict__ cos_tab, + const __nv_bfloat16 *__restrict__ sin_tab, + __nv_bfloat16 *__restrict__ q_out, + __nv_bfloat16 *__restrict__ k_out, + const int n_heads, + const int seq_len, + const __nv_bfloat16 inv_n, // bf16(1.0 / HEAD_DIM) -- affine mul + const __nv_bfloat16 zero, // bf16(0.0) -- affine add + const __nv_bfloat16 one, // bf16(1.0) -- affine mul + const __nv_bfloat16 eps, // bf16(rms_norm_eps) -- affine add + const int pos_offset) { + constexpr int NOPE = HEAD_DIM - ROPE_DIM; + constexpr int HALF_ROPE = ROPE_DIM / 2; + + const int row = blockIdx.x; // 0..n_heads-1 => Q head; n_heads => K + const int bt = blockIdx.y; // flattened (b, t) + const int t = bt % seq_len; + const int b = bt / seq_len; + const int tid = threadIdx.x; + const bool is_k = (row == n_heads); + + const __nv_bfloat16 *src; + __nv_bfloat16 *dst; + if (is_k) { + // K: [B, T, 1, D] -> [B, 1, T, D] + src = k_in + (size_t)bt * HEAD_DIM; + dst = k_out + (size_t)((size_t)b * seq_len + t) * HEAD_DIM; + } else { + // Q: [B, T, H, D] -> [B, H, T, D] (the materialised transpose) + src = q_in + ((size_t)bt * n_heads + row) * HEAD_DIM; + dst = q_out + (((size_t)b * n_heads + row) * seq_len + t) * HEAD_DIM; + } + + // -- 1. Q only: the RMS statistic --------------------------------------- + // + // candle: `q.sqr()?.mean_keepdim(-1)?.affine(1.0, eps)?.recip()?.sqrt()?` + // then `q.broadcast_mul(&inv_rms)`. K does NOT get this -- it has already + // been through `kv_norm` (a weighted RMSNorm module) before reaching here. + // Initialised to exactly 1.0 only so it is never read uninitialised on the + // K row; every use is guarded by `!is_k`, because K has already been + // through `kv_norm` and must NOT be scaled again. + __nv_bfloat16 inv_rms = one; + if (!is_k) { + __shared__ __nv_bfloat16 shr[HEAD_DIM]; + // `usqr_bf16` is `x*x`, rounded to bf16 BEFORE fast_sum reads it, so + // the multiply and the accumulate are two separate roundings here too. + __nv_bfloat16 x = src[tid]; + shr[tid] = x * x; + // candle fast_sum's tree, verbatim: __syncthreads() at the TOP of the + // body, and the halving order it implies. + for (int s = HEAD_DIM / 2; s > 0; s >>= 1) { + __syncthreads(); + if (tid < s) shr[tid] += shr[tid + s]; + } + // candle stops here (only tid 0 reads shr[0]); every thread needs the + // statistic to scale its own element, so one more barrier is added. + // A barrier changes no arithmetic. + __syncthreads(); + __nv_bfloat16 sum = shr[0]; + + // `mean_keepdim` == `sum_keepdim` then `* (1/n)`, which candle lowers + // to `affine(mul = 1/n, add = 0)`: `x * mul + add`. + __nv_bfloat16 mean = sum * inv_n + zero; + // `.affine(1.0, eps)` + __nv_bfloat16 denom = mean * one + eps; + inv_rms = qk_sqrt(qk_recip(denom)); + } + + // -- 2. NoPE dims: scale (Q) or copy (K) -------------------------------- + // + // This is the `cat([nope, pe])` low half. candle materialises it with a + // copy2d out of the normalised tensor; here it is written straight to the + // transposed destination. + for (int i = tid; i < NOPE; i += HEAD_DIM) { + dst[i] = is_k ? src[i] : (__nv_bfloat16)(src[i] * inv_rms); + } + + // -- 3. RoPE on the trailing ROPE_DIM dims ------------------------------ + // + // candle `ropei` (candle-kernels/src/reduce.cu): + // dst[2i] = src[2i] * c - src[2i+1] * s + // dst[2i+1] = src[2i] * s + src[2i+1] * c + // all in T. `rope_idx` there resolves to `t * (ROPE_DIM/2) + j` against the + // narrowed table, i.e. `(pos_offset + t) * (ROPE_DIM/2) + j` against the + // full one. + const __nv_bfloat16 *cos_row = cos_tab + (size_t)(pos_offset + t) * HALF_ROPE; + const __nv_bfloat16 *sin_row = sin_tab + (size_t)(pos_offset + t) * HALF_ROPE; + for (int j = tid; j < HALF_ROPE; j += HEAD_DIM) { + // The rotation reads the NORMALISED value, matching candle's order + // (broadcast_mul happens before apply_rope_inplace splits the tensor). + __nv_bfloat16 a = src[NOPE + 2 * j]; + __nv_bfloat16 bb = src[NOPE + 2 * j + 1]; + if (!is_k) { + a = a * inv_rms; + bb = bb * inv_rms; + } + const __nv_bfloat16 c = cos_row[j]; + const __nv_bfloat16 s = sin_row[j]; + // candle's `ropei` spells this `src[2i]*c - src[2i+1]*s`, and its PTX + // is contracted by the driver JIT so that the FIRST product is fused + // into the add/sub (one rounding) while the SECOND is rounded on its + // own. Written out explicitly because the operator form compiles to + // inline PTX here, which nvcc cannot contract at any `--fmad` setting. + // See the bit-identity contract at the top of this file for the + // measurement that pins this down. + const __nv_bfloat16 bs = bb * s; // round(b*s) + const __nv_bfloat16 bc = bb * c; // round(b*c) + dst[NOPE + 2 * j] = __hfma(a, c, __hneg(bs)); // round(a*c - round(b*s)) + dst[NOPE + 2 * j + 1] = __hfma(a, s, bc); // round(a*s + round(b*c)) + } +} + +extern "C" { + +// Returns 0 if the kernel was launched, non-zero if this (dtype, head_dim, +// rope_dim) combination is outside the specialised set. A non-zero return means +// the caller MUST fall back to the eager chain: producing a wrong answer from a +// generic slow path would be worse than the launches it saves. +int arc_qk_norm_rope_bf16_v2( + const void *q_in, + const void *k_in, + const void *cos_tab, + const void *sin_tab, + void *q_out, + void *k_out, + int n_heads, + int batch, + int seq_len, + int head_dim, + int rope_dim, + int pos_offset, + int dtype, + // Widened from uint16_t to a full word each. The bf16 bit patterns are + // carried in the low 16 bits. Sub-word arguments packed behind five other + // stack arguments are exactly the kind of boundary that is cheaper to + // remove than to reason about. + unsigned int inv_n_bits, + unsigned int zero_bits, + unsigned int one_bits, + unsigned int eps_bits, + long long stream) { + if (dtype != QK_DTYPE_BF16) return 1; + + __nv_bfloat16 inv_n, zero, one, eps; + unsigned short b; + b = (unsigned short)(inv_n_bits & 0xFFFFu); memcpy(&inv_n, &b, sizeof(b)); + b = (unsigned short)(zero_bits & 0xFFFFu); memcpy(&zero, &b, sizeof(b)); + b = (unsigned short)(one_bits & 0xFFFFu); memcpy(&one, &b, sizeof(b)); + b = (unsigned short)(eps_bits & 0xFFFFu); memcpy(&eps, &b, sizeof(b)); + + // One-shot probe: print what actually crossed the FFI boundary, once, so a + // wrong value is OBSERVED rather than inferred from output symptoms. + static int probed = 0; + if (!probed) { + probed = 1; + fprintf(stderr, + "[qk_norm_rope FFI] n_heads=%d batch=%d seq_len=%d head_dim=%d " + "rope_dim=%d pos_offset=%d dtype=%d inv_n=0x%04x zero=0x%04x " + "one=0x%04x eps=0x%04x\n", + n_heads, batch, seq_len, head_dim, rope_dim, pos_offset, dtype, + inv_n_bits & 0xFFFFu, zero_bits & 0xFFFFu, one_bits & 0xFFFFu, + eps_bits & 0xFFFFu); + fflush(stderr); + } + + dim3 grid((unsigned)(n_heads + 1), (unsigned)(batch * seq_len), 1); + cudaStream_t s = (cudaStream_t)stream; + +#define QK_LAUNCH(HD, RD) \ + qk_norm_rope_kernel<<>>( \ + (const __nv_bfloat16 *)q_in, (const __nv_bfloat16 *)k_in, \ + (const __nv_bfloat16 *)cos_tab, (const __nv_bfloat16 *)sin_tab, \ + (__nv_bfloat16 *)q_out, (__nv_bfloat16 *)k_out, n_heads, seq_len, \ + inv_n, zero, one, eps, pos_offset) + + // Only shapes whose row length IS candle's fast_sum block width are + // specialised -- see the bit-identity contract, point (1). V4 Flash is + // (512, 64); the others are here so a config change fails loudly at the + // dispatch rather than silently rounding differently. + if (head_dim == 512 && rope_dim == 64) { + QK_LAUNCH(512, 64); + } else if (head_dim == 256 && rope_dim == 64) { + QK_LAUNCH(256, 64); + } else if (head_dim == 192 && rope_dim == 64) { + // 192 is not a power of two: candle's block width would be 256, not + // 192, so the tree shape would differ. Refuse. + return 2; + } else { + return 2; + } +#undef QK_LAUNCH + + return 0; +} + +} // extern "C" diff --git a/mistralrs-core/src/cuda/qk_norm_rope.rs b/mistralrs-core/src/cuda/qk_norm_rope.rs new file mode 100644 index 000000000..68314ec23 --- /dev/null +++ b/mistralrs-core/src/cuda/qk_norm_rope.rs @@ -0,0 +1,445 @@ +//! Parent system: ArcInfer / ArcAttention +//! +//! Rust side of `cuda/qk_norm_rope.cu` — the fused DeepSeek-V4 Q/K +//! pre-attention block. +//! +//! One kernel replaces the sixteen candle launches that +//! [`crate::models::deepseek4::Attention::forward`] currently spends between +//! the Q/KV projections and attention: the head transpose, the per-head Q +//! RMS-normalisation, the RoPE split/rotate/`cat`, and the two +//! `.contiguous()` calls that materialise the intermediates. Ten of the +//! sixteen are pure data movement (`ucopy_bf16` / `copy2d_bf16`), which is the +//! block this file exists to remove. +//! +//! The kernel is **bit-identical** to the chain it replaces, not an +//! approximation. The contract, and the three reasons it is achievable at all, +//! are documented at the top of `qk_norm_rope.cu`. The short version: candle's +//! `fast_sum` reduction tree shape is copied rather than chosen, bf16 +//! arithmetic in CUDA is inline PTX and therefore immune to `--fmad` and +//! `--use_fast_math`, and the `affine` scalars are narrowed to bf16 host-side +//! by the same `half::bf16` code candle uses. +//! +//! Two switches, both read once and cached (`std::env::var` takes a global +//! lock and this is consulted on every layer of every step): +//! +//! - `ARC_QK_FUSED=0` restores the eager chain, so the A/B runs from ONE +//! binary and cannot be contaminated by a rebuild. +//! - `ARC_QK_VERIFY=1` runs BOTH paths at every layer and asserts the outputs +//! are bit-identical, with a negative control that proves the comparator is +//! live. See [`verify_enabled`]. + +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; + +/// Number of times the fused kernel has actually run. +/// +/// This exists because "the benchmark got faster" and "the kernel ran" are +/// different claims, and a fused path that silently declines every shape would +/// produce a perfectly green no-op. The A/B harness reads +/// [`engaged_count`] and treats zero as an environment failure, not a result. +static ENGAGED: AtomicU64 = AtomicU64::new(0); +/// Number of times the shape/dtype gate refused and the eager chain ran. +static DECLINED: AtomicU64 = AtomicU64::new(0); +static LOGGED_FIRST: AtomicBool = AtomicBool::new(false); + +/// `ARC_QK_FUSED=0` disables the fused Q/K kernel and restores the eager +/// candle chain. Any other value (or unset) keeps it on. +pub fn fused_enabled() -> bool { + use std::sync::OnceLock; + static ENABLED: OnceLock = OnceLock::new(); + *ENABLED.get_or_init(|| !matches!(std::env::var("ARC_QK_FUSED").as_deref(), Ok("0"))) +} + +/// `ARC_QK_VERIFY=1` runs the eager chain alongside the fused kernel at every +/// layer and compares the raw output bytes. +/// +/// The comparison is on bf16 BIT PATTERNS, not a tolerance — an earlier +/// bf16-level assertion elsewhere in this codebase was vacuous because eight +/// mantissa bits swallowed the difference it claimed to detect. Bit equality +/// cannot swallow anything, and [`verify_pair`] additionally runs a negative +/// control (one element poisoned by a single ULP) that MUST be flagged, +/// so a comparator that had silently degenerated into `true` is caught. +pub fn verify_enabled() -> bool { + use std::sync::OnceLock; + static ENABLED: OnceLock = OnceLock::new(); + *ENABLED.get_or_init(|| matches!(std::env::var("ARC_QK_VERIFY").as_deref(), Ok("1"))) +} + +/// `(engaged, declined)` — fused launches and shape-gate refusals so far. +pub fn engaged_count() -> (u64, u64) { + (ENGAGED.load(Ordering::Relaxed), DECLINED.load(Ordering::Relaxed)) +} + +pub(crate) fn note_engaged(head_dim: usize, rope_dim: usize, n_heads: usize, dtype: &str) { + let n = ENGAGED.fetch_add(1, Ordering::Relaxed); + if !LOGGED_FIRST.swap(true, Ordering::Relaxed) { + tracing::info!( + "ArcAttention: fused qk_norm_rope ENGAGED (head_dim={head_dim}, rope_dim={rope_dim}, \ + n_heads={n_heads}, dtype={dtype}) — replacing 16 candle launches per layer with 1" + ); + } + // Periodic proof of life for long runs; cheap (one relaxed load). + if n > 0 && n % 100_000 == 0 { + tracing::info!("ArcAttention: fused qk_norm_rope engaged {n} times"); + } +} + +pub(crate) fn note_declined(reason: &str) { + let n = DECLINED.fetch_add(1, Ordering::Relaxed); + if n == 0 { + tracing::warn!( + "ArcAttention: fused qk_norm_rope DECLINED ({reason}) — falling back to the eager \ + chain. This is correct but slow; the launch saving is not being taken." + ); + } +} + +#[cfg(feature = "cuda")] +mod cuda_impl { + use candle_core as candle; + use candle_core::{DType, Result, Tensor}; + + /// Pull a contiguous BF16 CUDA tensor's device pointer. + fn bf16_ptr(t: &Tensor, what: &str) -> Result<*const std::ffi::c_void> { + use candle_core::cuda_backend::cudarc::driver::DevicePtr; + if t.dtype() != DType::BF16 { + candle::bail!("qk_norm_rope: {what} must be BF16, got {:?}", t.dtype()); + } + if !t.is_contiguous() { + candle::bail!("qk_norm_rope: {what} must be contiguous"); + } + let (s, l) = t.storage_and_layout(); + let cuda = match &*s { + candle::Storage::Cuda(c) => c, + _ => candle::bail!("qk_norm_rope: {what} must be on CUDA"), + }; + // The (ptr, guard) pair borrows the slice, so bind before returning — + // same shape as `hc_fused::any_ptr`. + let sl = cuda.as_cuda_slice::()?; + let ptr = sl.slice(l.start_offset()..).device_ptr(sl.stream()).0; + Ok(ptr as *const std::ffi::c_void) + } + + /// Fused head-transpose + Q RMS-norm + RoPE + NoPE/PE recombination. + /// + /// - `q_in` `[B, T, H * D]` or `[B, T, H, D]` BF16 contiguous — the raw + /// q_proj output, **not** transposed. The transpose is the kernel's + /// output addressing. + /// - `k_in` `[B, T, D]` or `[B, T, 1, D]` BF16 contiguous — post-`kv_norm`. + /// - `cos` / `sin` `[max_pos, D_rope / 2]` BF16 contiguous — the FULL + /// tables; `pos_offset` replaces candle's `narrow`. + /// + /// Returns `(q_out [B, H, T, D], k_out [B, 1, T, D])`. + #[allow(clippy::too_many_arguments)] + pub fn qk_norm_rope_cuda( + q_in: &Tensor, + k_in: &Tensor, + cos: &Tensor, + sin: &Tensor, + n_heads: usize, + batch: usize, + seq_len: usize, + head_dim: usize, + rope_dim: usize, + eps: f64, + pos_offset: usize, + ) -> Result<(Tensor, Tensor)> { + use candle_core::cuda_backend::cudarc::driver::DevicePtr; + + let q_ptr = bf16_ptr(q_in, "q_in")?; + let k_ptr = bf16_ptr(k_in, "k_in")?; + let cos_ptr = bf16_ptr(cos, "cos")?; + let sin_ptr = bf16_ptr(sin, "sin")?; + + let dev = q_in.device().as_cuda_device()?; + let stream = dev.cuda_stream().cu_stream() as i64; + + let q_elems = batch * n_heads * seq_len * head_dim; + let k_elems = batch * seq_len * head_dim; + + // candle's `affine(mul, add)` narrows its f64 arguments to T before the + // kernel sees them. Do the same narrowing, with the same code, so no + // rounding decision is taken twice. + let inv_n = u32::from(half::bf16::from_f64(1.0f64 / head_dim as f64).to_bits()); + let zero = u32::from(half::bf16::from_f64(0.0).to_bits()); + let one = u32::from(half::bf16::from_f64(1.0).to_bits()); + let eps_b = u32::from(half::bf16::from_f64(eps).to_bits()); + + let q_buf = unsafe { dev.alloc::(q_elems) }?; + let k_buf = unsafe { dev.alloc::(k_elems) }?; + + let rc = unsafe { + crate::cuda::ffi::arc_qk_norm_rope_bf16_v2( + q_ptr, + k_ptr, + cos_ptr, + sin_ptr, + q_buf.device_ptr(q_buf.stream()).0 as *mut std::ffi::c_void, + k_buf.device_ptr(k_buf.stream()).0 as *mut std::ffi::c_void, + n_heads as i32, + batch as i32, + seq_len as i32, + head_dim as i32, + rope_dim as i32, + pos_offset as i32, + 1, // QK_DTYPE_BF16 + inv_n, + zero, + one, + eps_b, + stream, + ) + }; + if rc != 0 { + // The host-side gate in `deepseek4.rs` should have caught this, so + // reaching here is a bug rather than an unsupported config. + candle::bail!( + "qk_norm_rope: kernel dispatch refused (rc={rc}) for head_dim={head_dim} \ + rope_dim={rope_dim} — host gate and device gate disagree" + ); + } + + let q_st = candle::CudaStorage::wrap_cuda_slice(q_buf, dev.clone()); + let k_st = candle::CudaStorage::wrap_cuda_slice(k_buf, dev.clone()); + let q_out = Tensor::from(( + candle::Storage::Cuda(q_st), + candle_core::Shape::from((batch, n_heads, seq_len, head_dim)), + )); + let k_out = Tensor::from(( + candle::Storage::Cuda(k_st), + candle_core::Shape::from((batch, 1usize, seq_len, head_dim)), + )); + Ok((q_out, k_out)) + } +} + +#[cfg(feature = "cuda")] +pub use cuda_impl::qk_norm_rope_cuda; + +#[cfg(not(feature = "cuda"))] +#[allow(clippy::too_many_arguments)] +pub fn qk_norm_rope_cuda( + _q_in: &candle_core::Tensor, + _k_in: &candle_core::Tensor, + _cos: &candle_core::Tensor, + _sin: &candle_core::Tensor, + _n_heads: usize, + _batch: usize, + _seq_len: usize, + _head_dim: usize, + _rope_dim: usize, + _eps: f64, + _pos_offset: usize, +) -> candle_core::Result<(candle_core::Tensor, candle_core::Tensor)> { + candle_core::bail!("qk_norm_rope_cuda requires the cuda feature") +} + +/// Count bit-pattern mismatches between two bf16 slices. +/// +/// The ONE comparison routine used by both the real check and its negative +/// control, so the control provably exercises the same code. +fn bit_mismatches(a: &[half::bf16], b: &[half::bf16]) -> (usize, Option<(usize, u16, u16)>) { + let mut n = 0usize; + let mut first = None; + for (i, (x, y)) in a.iter().zip(b.iter()).enumerate() { + if x.to_bits() != y.to_bits() { + n += 1; + if first.is_none() { + first = Some((i, x.to_bits(), y.to_bits())); + } + } + } + (n, first) +} + +/// Bit-identity check with a built-in negative control. +/// +/// `fused` and `eager` must be byte-for-byte equal. The control runs +/// [`bit_mismatches`] — the very same routine — against a copy of `eager` with +/// one element perturbed by a single bf16 ULP. If that run does NOT report +/// exactly one difference, the comparator is not doing its job and the check +/// is declared VOID rather than passed. +/// +/// Comparison is on bf16 BIT PATTERNS, never a tolerance. An earlier +/// bf16-level assertion in this codebase was vacuous precisely because eight +/// mantissa bits swallowed the difference it claimed to detect; bit equality +/// cannot swallow anything. +/// +/// Returns `Ok(())` only on a proven match — a verify leg must be able to exit +/// non-zero. +pub fn verify_pair( + name: &str, + fused: &candle_core::Tensor, + eager: &candle_core::Tensor, +) -> candle_core::Result<()> { + if fused.shape() != eager.shape() { + candle_core::bail!( + "qk_norm_rope VERIFY {name}: shape {:?} != {:?}", + fused.shape(), + eager.shape() + ); + } + if fused.dtype() != candle_core::DType::BF16 || eager.dtype() != candle_core::DType::BF16 { + candle_core::bail!( + "qk_norm_rope VERIFY {name}: expected BF16/BF16, got {:?}/{:?}", + fused.dtype(), + eager.dtype() + ); + } + let f: Vec = fused.flatten_all()?.to_vec1()?; + let e: Vec = eager.flatten_all()?.to_vec1()?; + if f.len() != e.len() { + candle_core::bail!("qk_norm_rope VERIFY {name}: len {} != {}", f.len(), e.len()); + } + if f.is_empty() { + candle_core::bail!("qk_norm_rope VERIFY {name}: empty tensor — nothing was compared"); + } + + // -- negative control, run FIRST so a degenerate comparator cannot pass -- + let poison_idx = f.len() / 2; + let mut poisoned = e.clone(); + poisoned[poison_idx] = half::bf16::from_bits(e[poison_idx].to_bits() ^ 1); + let (ctrl_n, _) = bit_mismatches(&poisoned, &e); + if ctrl_n != 1 { + candle_core::bail!( + "qk_norm_rope VERIFY {name}: NEGATIVE CONTROL FAILED — a 1-ULP poison at index \ + {poison_idx} produced {ctrl_n} detected mismatches, expected exactly 1. The \ + comparator is not measuring what it claims; declaring VOID rather than passed." + ); + } + + // -- the real comparison, same routine -- + let (n, first) = bit_mismatches(&f, &e); + if n != 0 { + let (i, fb, eb) = first.expect("mismatch count > 0 implies a first mismatch"); + // WHERE a mismatch falls inside the head vector separates the two + // failure modes that look alike from a bare count: a wrong RMS + // statistic scatters over all dims, while a wrong RoPE table row hits + // only the low-j (high-frequency) pairs, because high-j pairs have + // cos ~ 1 / sin ~ 0 at any nearby position and round to the same bf16. + let hd = *fused.dims().last().unwrap_or(&1); + let mut nope_bad = 0usize; + let mut pairs: Vec = Vec::new(); + let rope = hd.min(64); + let nope = hd - rope; + for (idx, (x, y)) in f.iter().zip(e.iter()).enumerate() { + if x.to_bits() == y.to_bits() { + continue; + } + let d = idx % hd; + if d < nope { + nope_bad += 1; + } else { + let j = (d - nope) / 2; + if !pairs.contains(&j) { + pairs.push(j); + } + } + } + pairs.sort_unstable(); + candle_core::bail!( + "qk_norm_rope VERIFY {name}: {n}/{} elements differ; first at {i}: \ + fused=0x{fb:04x} eager=0x{eb:04x}; head_dim={hd} nope_mismatches={nope_bad} \ + rope_pairs_hit={pairs:?}", + f.len() + ); + } + VERIFIED.fetch_add(f.len() as u64, Ordering::Relaxed); + Ok(()) +} + +/// `ARC_QK_DIAG=1` runs [`diagnose_row`] on the first fused call. +pub fn diag_enabled() -> bool { + use std::sync::OnceLock; + static ENABLED: OnceLock = OnceLock::new(); + *ENABLED.get_or_init(|| matches!(std::env::var("ARC_QK_DIAG").as_deref(), Ok("1"))) +} + +static DIAGGED: AtomicBool = AtomicBool::new(false); + +/// Recover, from candle's own output, WHICH row of the RoPE table it rotated +/// by — and report whether that is the row this kernel reads. +/// +/// `q_normed` is the eager normalised Q `[B, H, T, D]`; `q_eager` is the eager +/// post-RoPE Q of the same shape. Only head 0, t 0, pair 0 is needed: it is the +/// highest-frequency pair, so it pins the position most sharply. +pub fn diagnose_row( + q_normed: &candle_core::Tensor, + q_eager: &candle_core::Tensor, + q_fused: &candle_core::Tensor, + cos: &candle_core::Tensor, + sin: &candle_core::Tensor, + head_dim: usize, + rope_dim: usize, + pos_offset: usize, + seq_len: usize, +) -> candle_core::Result<()> { + if DIAGGED.swap(true, Ordering::Relaxed) { + return Ok(()); + } + tracing::error!("qk_norm_rope DIAG: pos_offset={pos_offset} seq_len={seq_len} head_dim={head_dim} rope_dim={rope_dim}"); + let nope = head_dim - rope_dim; + let n: Vec = q_normed.flatten_all()?.to_dtype(candle_core::DType::F32)?.to_vec1()?; + let o: Vec = q_eager.flatten_all()?.to_dtype(candle_core::DType::F32)?.to_vec1()?; + let fu: Vec = q_fused.flatten_all()?.to_dtype(candle_core::DType::F32)?.to_vec1()?; + let a = n[nope] as f64; + let b = n[nope + 1] as f64; + let o0 = o[nope] as f64; + let o1 = o[nope + 1] as f64; + let den = a * a + b * b; + if den == 0.0 { + tracing::warn!("qk_norm_rope DIAG: pair 0 is zero; cannot solve for the row"); + return Ok(()); + } + let c_used = (a * o0 + b * o1) / den; + let s_used = (a * o1 - b * o0) / den; + + let (rows, cols) = cos.dims2()?; + let half = rope_dim / 2; + let scan = rows.min(pos_offset + 4096); + let cv: Vec = cos.narrow(0, 0, scan)?.flatten_all()?.to_dtype(candle_core::DType::F32)?.to_vec1()?; + let sv: Vec = sin.narrow(0, 0, scan)?.flatten_all()?.to_dtype(candle_core::DType::F32)?.to_vec1()?; + let mut best = (f64::INFINITY, usize::MAX); + for r in 0..scan { + let dc = cv[r * cols] as f64 - c_used; + let ds = sv[r * cols] as f64 - s_used; + let d = dc * dc + ds * ds; + if d < best.0 { + best = (d, r); + } + } + let mine = pos_offset; + let cb: Vec = cos.narrow(0, mine, 1)?.flatten_all()?.to_vec1()?; + let sb: Vec = sin.narrow(0, mine, 1)?.flatten_all()?.to_vec1()?; + tracing::error!( + "qk_norm_rope DIAG bits: kernel row {mine} pair0 cos=0x{:04x} sin=0x{:04x}; \ + pairs 0..4 cos=[{:04x},{:04x},{:04x},{:04x}] sin=[{:04x},{:04x},{:04x},{:04x}]", + cb[0].to_bits(), sb[0].to_bits(), + cb[0].to_bits(), cb[1].to_bits(), cb[2].to_bits(), cb[3].to_bits(), + sb[0].to_bits(), sb[1].to_bits(), sb[2].to_bits(), sb[3].to_bits() + ); + tracing::error!( + "qk_norm_rope DIAG: table is [{rows}, {cols}], half={half}, seq_len={seq_len}. \ + candle rotated pair 0 by (cos={c_used:.6}, sin={s_used:.6}); the NEAREST table row is \ + {} (err={:.3e}); the kernel reads row {mine} whose pair 0 is (cos={:.6}, sin={:.6}). \ + ROW_USED_BY_CANDLE={} ROW_READ_BY_KERNEL={mine}", + best.1, + best.0.sqrt(), + cv[mine * cols], + sv[mine * cols], + best.1 + ); + tracing::error!( + "qk_norm_rope DIAG pair0: normed a={a:.6} b={b:.6} | eager out=({o0:.6}, {o1:.6}) | fused out=({:.6}, {:.6}) | nope[0] eager={:.6} fused={:.6}", + fu[nope], fu[nope + 1], o[0], fu[0] + ); + Ok(()) +} + +/// Total bf16 elements proven bit-identical so far. Reported by the verify leg +/// so "0 mismatches" can never be confused with "nothing was compared". +static VERIFIED: AtomicU64 = AtomicU64::new(0); + +/// Elements proven bit-identical so far. +pub fn verified_elements() -> u64 { + VERIFIED.load(Ordering::Relaxed) +} diff --git a/mistralrs-core/src/layers.rs b/mistralrs-core/src/layers.rs index 94fd20401..39eefbad2 100644 --- a/mistralrs-core/src/layers.rs +++ b/mistralrs-core/src/layers.rs @@ -1607,6 +1607,19 @@ impl DeepSeekV2RotaryEmbedding { } } + /// The full `[max_position_embeddings, rope_dim / 2]` cosine table. + /// + /// Exposed for `cuda::qk_norm_rope`, which addresses the table with a base + /// offset instead of materialising the `narrow` view `forward` takes. + pub fn cos(&self) -> &Tensor { + &self.cos + } + + /// The full `[max_position_embeddings, rope_dim / 2]` sine table. See [`Self::cos`]. + pub fn sin(&self) -> &Tensor { + &self.sin + } + pub fn forward( &self, q: &Tensor, diff --git a/mistralrs-core/src/models/deepseek4.rs b/mistralrs-core/src/models/deepseek4.rs index f18e6a093..e2f6f6336 100644 --- a/mistralrs-core/src/models/deepseek4.rs +++ b/mistralrs-core/src/models/deepseek4.rs @@ -1235,6 +1235,99 @@ impl Attention { }) } + /// Decide whether the fused `cuda::qk_norm_rope` kernel may replace the + /// eager transpose + Q RMS-norm + RoPE chain for this call. + /// + /// Returns `Some(pos_offset)` when every precondition for BIT-IDENTITY + /// holds, `None` otherwise. Each rejection is a real difference in what the + /// kernel would have to compute, not a convenience: + /// + /// * **BF16 only.** In f32 the RoPE expression is FMA-contractible, and + /// candle-kernels builds with the default `-fmad=true` while this + /// kernel's IEEE builder uses `--fmad=false`; the two would legitimately + /// disagree in the last bit. bf16 operators are inline PTX and immune. + /// * **`head_dim` must equal candle's `fast_sum` block width** (i.e. be a + /// power of two ≤ 1024), or the reduction tree — and therefore the + /// rounding of the RMS statistic — has a different shape. + /// * **One shared position offset.** With per-sequence offsets candle takes + /// its `Tensor::cat` path, which the kernel does not model. + /// * **A single KV head**, which is what makes K one extra row rather than + /// a second layout. + /// + /// A `None` here is logged once (see `note_declined`) so a permanently + /// disengaged fast path cannot masquerade as a working one. + fn fused_qk_gate( + &self, + q_proj_out: &Tensor, + kv_normed: &Tensor, + seqlen_offsets: &[usize], + ) -> Option { + use crate::cuda::qk_norm_rope as qk; + if !cfg!(feature = "cuda") || !qk::fused_enabled() { + return None; + } + macro_rules! decline { + ($why:expr) => {{ + qk::note_declined($why); + return None; + }}; + } + if !q_proj_out.device().is_cuda() { + decline!("not a CUDA device"); + } + if q_proj_out.dtype() != candle_core::DType::BF16 + || kv_normed.dtype() != candle_core::DType::BF16 + { + decline!("Q/KV are not BF16"); + } + let cos = self.rotary_emb.cos(); + let sin = self.rotary_emb.sin(); + if cos.dtype() != candle_core::DType::BF16 || sin.dtype() != candle_core::DType::BF16 { + decline!("RoPE tables are not BF16"); + } + if seqlen_offsets.len() != 1 { + decline!("per-sequence position offsets"); + } + if self.num_kv_heads != 1 { + decline!("more than one KV head"); + } + let head_dim = self.cfg.head_dim; + let rope_dim = self.cfg.qk_rope_head_dim; + // candle's FastReduce block width is `min(1024, n).next_power_of_two()`; + // the kernel's tree only matches when that equals the row length. + if head_dim > 1024 || !head_dim.is_power_of_two() { + decline!("head_dim is not candle's fast_sum block width"); + } + if rope_dim == 0 || rope_dim % 2 != 0 || rope_dim > head_dim { + decline!("unusable rope_dim"); + } + if !q_proj_out.is_contiguous() || !kv_normed.is_contiguous() { + decline!("Q/KV projection output is not contiguous"); + } + if !cos.is_contiguous() || !sin.is_contiguous() { + decline!("RoPE tables are not contiguous"); + } + // Shapes: q_proj_out is [B, T, n_heads * head_dim], kv_normed is + // [B, T, head_dim]. + let (qb, qt, qh) = match q_proj_out.dims3() { + Ok(d) => d, + Err(_) => decline!("q_proj output is not 3-D"), + }; + if qh != self.num_attention_heads * head_dim { + decline!("q_proj output width is not n_heads * head_dim"); + } + match kv_normed.dims3() { + Ok((kb, kt, kh)) if kb == qb && kt == qt && kh == head_dim => {} + _ => decline!("kv_norm output shape does not match [B, T, head_dim]"), + } + let pos_offset = seqlen_offsets[0]; + match cos.dims2() { + Ok((rows, cols)) if cols == rope_dim / 2 && rows >= pos_offset + qt => {} + _ => decline!("RoPE table does not cover [pos_offset, pos_offset + T)"), + } + Some(pos_offset) + } + /// Apply RoPE in-place to the last `qk_rope_head_dim` dims of each /// Q-head and K-head's `head_dim`-vector. The first /// `qk_nope_head_dim` dims are left untouched. Audit §0 + §3. @@ -1451,36 +1544,21 @@ impl Attention { drop(_prof_comp); // 1. Q projection (LoRA). [B, T, hidden] → [B, T, n_heads*head_dim] // → reshape to [B, n_heads, T, head_dim]. Audit §3. - let q = { + let q_proj_out = { let _s = arc_profiler::device_span("q_proj"); timed_mla(0, &mdev, || self.q.forward(xs))? }; - v4_nan_dbg(&q, "attn.q_proj"); - let q = q - .reshape((bs, seq_len, self.num_attention_heads, head_dim))? - .transpose(1, 2)? - .contiguous()?; - // V4: per-head RMS-normalize Q over head_dim before RoPE. Reference - // inference/model.py:498 `q *= rsqrt(q.square().mean(-1)+eps)`. This is - // SEPARATE from q_norm (which normalizes q_lora_rank inside self.q). - // Missing it leaves Q ~30x too small -> near-uniform attention scores - // -> the model cannot attend -> word-salad output (RUN-161). - let q = { - let _s = arc_profiler::device_span("q_rmsnorm"); - let inv_rms = q - .sqr()? - .mean_keepdim(candle_core::D::Minus1)? - .affine(1.0, self.cfg.rms_norm_eps)? - .recip()? - .sqrt()?; - q.broadcast_mul(&inv_rms)? - }; - v4_stat_dbg(&q, "attn.q_normed"); - v4_trace_dump(self.dbg_layer_idx, &q, "10_q_normed"); + v4_nan_dbg(&q_proj_out, "attn.q_proj"); // 2. K/V projection: single fused wkv. [B, T, hidden] → [B, T, // head_dim] → kv_norm → reshape to [B, num_kv_heads=1, T, // head_dim]. Audit §0 + §3. + // + // wave44: this block used to sit AFTER the Q RMS-norm. It is hoisted + // above it so the fused Q/K kernel below can take the Q rows and the K + // row in ONE launch. The two projections read only `xs` and neither + // reads the other's output, so no arithmetic moves — only the order in + // which two independent GEMMs are enqueued. let kv_raw = { let _s = arc_profiler::device_span("kv_proj"); timed_mla(1, &mdev, || self.wkv.forward_autocast(xs))? @@ -1491,16 +1569,119 @@ impl Attention { self.kv_norm.forward(&kv_raw)? }; v4_nan_dbg(&kv_normed, "attn.kv_norm"); - let k = kv_normed - .reshape((bs, seq_len, self.num_kv_heads, head_dim))? - .transpose(1, 2)? - .contiguous()?; - // 3. RoPE applied in-place to the last qk_rope_head_dim dims of - // each Q-head and K-head's head_dim-vector. Audit §0 + §3. - let (q, k) = { + // 3. Head transpose + per-head Q RMS-norm + RoPE + NoPE/PE recombine. + // + // The eager spelling below is SIXTEEN candle launches per layer, TEN of + // them pure data movement (`ucopy_bf16` / `copy2d_bf16`) materialising + // intermediates that only exist because the chain is expressed as + // separate ops: a transpose, a `broadcast_mul`, two `.contiguous()` + // narrows inside `rope_i`, and two two-input `cat`s. Their entire data + // footprint at decode is 64 KB of Q and 1 KB of K. + // + // `cuda::qk_norm_rope` does all of it in one launch of `n_heads + 1` + // blocks, bit-identically — the transpose becomes an output address, + // the split/`cat` becomes two write ranges of the same row, and the RMS + // statistic stays in a register. Toggle with `ARC_QK_FUSED=0`; prove + // with `ARC_QK_VERIFY=1`. + let eager_qk = |q_proj_out: &Tensor, kv_normed: &Tensor| -> Result<(Tensor, Tensor)> { + let q = q_proj_out + .reshape((bs, seq_len, self.num_attention_heads, head_dim))? + .transpose(1, 2)? + .contiguous()?; + // V4: per-head RMS-normalize Q over head_dim before RoPE. Reference + // inference/model.py:498 `q *= rsqrt(q.square().mean(-1)+eps)`. This is + // SEPARATE from q_norm (which normalizes q_lora_rank inside self.q). + // Missing it leaves Q ~30x too small -> near-uniform attention scores + // -> the model cannot attend -> word-salad output (RUN-161). + let q = { + let _s = arc_profiler::device_span("q_rmsnorm"); + let inv_rms = q + .sqr()? + .mean_keepdim(candle_core::D::Minus1)? + .affine(1.0, self.cfg.rms_norm_eps)? + .recip()? + .sqrt()?; + q.broadcast_mul(&inv_rms)? + }; + v4_stat_dbg(&q, "attn.q_normed"); + v4_trace_dump(self.dbg_layer_idx, &q, "10_q_normed"); + + let k = kv_normed + .reshape((bs, seq_len, self.num_kv_heads, head_dim))? + .transpose(1, 2)? + .contiguous()?; + + // RoPE applied in-place to the last qk_rope_head_dim dims of + // each Q-head and K-head's head_dim-vector. Audit §0 + §3. let _s = arc_profiler::device_span("rope"); - self.apply_rope_inplace(&q, &k, seqlen_offsets)? + self.apply_rope_inplace(&q, &k, seqlen_offsets) + }; + + let (q, k) = { + let _s = arc_profiler::device_span("qk_norm_rope"); + match self.fused_qk_gate(&q_proj_out, &kv_normed, seqlen_offsets) { + Some(pos_offset) => { + let (q_f, k_f) = crate::cuda::qk_norm_rope::qk_norm_rope_cuda( + &q_proj_out, + &kv_normed, + self.rotary_emb.cos(), + self.rotary_emb.sin(), + self.num_attention_heads, + bs, + seq_len, + head_dim, + self.cfg.qk_rope_head_dim, + self.cfg.rms_norm_eps as f64, + pos_offset, + )?; + crate::cuda::qk_norm_rope::note_engaged( + head_dim, + self.cfg.qk_rope_head_dim, + self.num_attention_heads, + "bf16", + ); + if crate::cuda::qk_norm_rope::verify_enabled() { + let (q_e, k_e) = eager_qk(&q_proj_out, &kv_normed)?; + // The eager NORMALISED Q (pre-RoPE), built on demand so + // the rotation candle applied can be solved for exactly. + let qn_for_diag = || -> Result { + let qn = q_proj_out + .reshape((bs, seq_len, self.num_attention_heads, head_dim))? + .transpose(1, 2)? + .contiguous()?; + let inv_rms = qn + .sqr()? + .mean_keepdim(candle_core::D::Minus1)? + .affine(1.0, self.cfg.rms_norm_eps)? + .recip()? + .sqrt()?; + qn.broadcast_mul(&inv_rms) + }; + // Run the row diagnostic on the FAILING call, not on + // whichever call happened to be first: position 0 is + // the identity rotation, so the first call agrees + // trivially and describes nothing. + if let Err(e) = crate::cuda::qk_norm_rope::verify_pair("q", &q_f, &q_e) { + crate::cuda::qk_norm_rope::diagnose_row( + &qn_for_diag()?, + &q_e, + &q_f, + self.rotary_emb.cos(), + self.rotary_emb.sin(), + head_dim, + self.cfg.qk_rope_head_dim, + seqlen_offsets[0], + seq_len, + )?; + return Err(e); + } + crate::cuda::qk_norm_rope::verify_pair("k", &k_f, &k_e)?; + } + (q_f, k_f) + } + None => eager_qk(&q_proj_out, &kv_normed)?, + } }; v4_nan_dbg(&q, "attn.q_rope"); v4_nan_dbg(&k, "attn.k_rope"); From 1f6ef9da911c997cb8939b8d7d69caa147cf9c65 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Thu, 20 Aug 2026 07:31:53 +0100 Subject: [PATCH 25/27] =?UTF-8?q?perf(ArcInfer/ArcMoE):=20close=20the=20F3?= =?UTF-8?q?2=20seam=20=E2=80=94=205=20sites,=20~950=20launches=20+=20258?= =?UTF-8?q?=20H2D=20copies=20per=20decode=20forward?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The router/expert region computes in F32 while the model is BF16. Every entry and exit paid a `cast_*` launch and every intermediate was a materialised F32 tensor. Measured on the H200 at b424f5cd9, per decode forward: cast_bf16_f32 356.7, cast_f32_bf16 221.1, cast_u32_f32 182.0, affine_f32 211.5, bminimum_f32 172.0, fast_sum_f32 98.8, bmul_f32 94.4, bmaximum_f32 86.0. `bmaximum_f32 = 86.0` and `bminimum_f32 = 172.0` are EXACT, and that is what identified the seam: 86 = 2 x 43 layers is the number of `swiglu_clamp` calls per forward (routed experts + shared expert), and each holds one `maximum` and two `minimum`s. Nothing else in the model produces a float `maximum`. A second cost never appears in a kernel histogram at all. candle's `minimum(f64)`/`maximum(f64)` build the bound as `Tensor::new(v, Cpu)?.to_dtype()?.to_device(cuda)?` on every call (`binary_op_scalar!`), i.e. a 4-byte host-to-device memcpy per clamp bound — three per `swiglu_clamp`, 258 per forward. Five sites, one master switch (`ARC_F32SEAM=0`) plus one per site so the whole change is A/B-able from a single binary: 1. Routed-expert clamped SwiGLU — 8 launches + 3 H2D -> 1 kernel, x43. 2. Shared-expert clamped SwiGLU — the CLAMP HALF ONLY, 5 launches + 3 H2D -> 1, x43. See below. 3. `experts.weighted_sum` — widen, weight, reduce over the 6-element expert axis, narrow: 4 launches -> 1, x43. 4. `gate.renormalize` — fast_sum, affine(+1e-20), bdiv, affine(*scale) over a [1, 6] tensor: 4 launches -> 1, x43. 5. `compressed_kv_from_rows` — `j * ratio` was laundered through float (`arange(u32) -> cast_u32_f32 -> affine -> cast_f32_u32`); `arange_step` produces the identical U32 range in one op. Plus one pure bug: `v4_collapse_dbg(&topk_idx.to_dtype(F32)..., ...)` put a cast in ARGUMENT position, so it ran on every MoE layer of every forward — 43 `cast_u32_f32` per token — even with `ARC_COLLAPSE` unset. `v4_collapse_dbg` already casts internally, behind the guard. WHY THE SHARED EXPERT ONLY GETS HALF. Its activation and multiply are `mistralrs_quant::fused_glu`, and mistralrs-quant/build.rs:281 compiles that translation unit with `--use_fast_math`: its `expf` is `__expf` (ex2.approx), its `/` is div.approx, and fast math also implies -ftz=true, which no intrinsic in a non-fast-math file can reproduce per-op. The full fusion reproduces candle's IEEE `usilu_f32` and would therefore NOT be bit-identical there. Rather than trade the contract for four launches, only the provable half is fused. `ARC_SEAM_AB=1` still evaluates the full fusion at that site and reports whether the two spellings agree — the measurement that would license the rest. BIT-IDENTITY, BY TRANSCRIPTION. Each kernel carries the op-by-op citation of the chain it replaces (candle `usilu_f32` = `x / (1 + expg(-x))`, `maxg`/`ming` = `fmaxf`/`fminf`, `clamp` = `maximum` then `minimum`, `Tensor +/* f64` = `affine` = a single `fmaf`, `sum` = the identity-padded `fast_sum` tree). The reduction order is the subtle one: `top_k = 6` reduces in a block of EIGHT, so two zero-padded lanes are part of the answer. `ARC_SEAM_AB=1` recomputes every eager chain beside its kernel and compares raw bits AT F32, BEFORE THE NARROWING — a BF16 comparison has 8 mantissa bits and would swallow exactly the reassociation errors this is meant to catch. `ARC_HC_AB_POISON=1` is the negative control that must make it fail. Engagement is reported, not assumed: every site logs its first ENGAGED and its first DECLINED (with the reason) and prints running totals, so "no DECLINED line" is a fact in the log rather than an inference. Five CPU tests pin what is checkable without a GPU: the padded reduction tree against `sinkhorn::reference::candle_tree_sum` (an independent transcription of the same candle kernel), the k=6 tree against the sequential sum it must NOT equal, `affine` as an fma rather than a multiply, the engagement counters moving in both directions, and the build wiring — including a tripwire that fires if mistralrs-quant ever drops `--use_fast_math`, which would make the shared expert's remaining four launches fusable. Co-Authored-By: Claude Opus 5 (1M context) --- mistralrs-core/src/cuda/ffi.rs | 54 +++ mistralrs-core/src/cuda/hc_fused.cu | 410 ++++++++++++++++++++++ mistralrs-core/src/cuda/hc_fused.rs | 463 ++++++++++++++++++++++++- mistralrs-core/src/layers.rs | 10 +- mistralrs-core/src/models/deepseek4.rs | 111 +++++- mistralrs-core/src/moe/experts.rs | 251 +++++++++++++- mistralrs-core/src/moe/mod.rs | 2 +- 7 files changed, 1279 insertions(+), 22 deletions(-) diff --git a/mistralrs-core/src/cuda/ffi.rs b/mistralrs-core/src/cuda/ffi.rs index ade2d091b..fbb0b9e43 100644 --- a/mistralrs-core/src/cuda/ffi.rs +++ b/mistralrs-core/src/cuda/ffi.rs @@ -440,6 +440,60 @@ extern "C" { dtype: i32, stream: i64, ) -> i32; + /// V4 clamped SwiGLU: `narrow(silu(min(gate, limit)) * clamp(up, -limit, limit))`. + /// Returns 0 on launch, non-zero if the `(in_dtype, out_dtype)` pair has no + /// specialisation — the caller must then use the eager chain. + #[allow(clippy::too_many_arguments)] + pub(crate) fn arc_seam_swiglu_clamp( + gate: *const c_void, + up: *const c_void, + out: *mut c_void, + limit: f32, + total: i64, + in_dtype: i32, + out_dtype: i32, + stream: i64, + ) -> i32; + /// The clamp half only, for the shared-expert site whose activation is a + /// fast-math `fused_glu` that a non-fast-math kernel cannot reproduce: + /// `gate_out = min(gate, limit)`, `up_out = clamp(up, -limit, limit)`. + #[allow(clippy::too_many_arguments)] + pub(crate) fn arc_seam_swiglu_clamp_split( + gate: *const c_void, + up: *const c_void, + gate_out: *mut c_void, + up_out: *mut c_void, + limit: f32, + total: i64, + in_dtype: i32, + stream: i64, + ) -> i32; + /// MoE expert combine: `out[n, h] = narrow(tree_sum_k(ys[n, j, h] * w[n, j]))`. + /// Same return contract as above; `k` outside the specialised set refuses. + #[allow(clippy::too_many_arguments)] + pub(crate) fn arc_seam_moe_weighted_sum( + ys: *const c_void, + w: *const c_void, + out: *mut c_void, + k: i32, + h: i32, + total: i64, + in_dtype: i32, + out_dtype: i32, + stream: i64, + ) -> i32; + /// MoE gate weight renormalise + scale. Same return contract as above. + #[allow(clippy::too_many_arguments)] + pub(crate) fn arc_seam_gate_renorm( + w: *const c_void, + out: *mut c_void, + eps: f32, + scale: f32, + do_renorm: i32, + k: i32, + n: i64, + stream: i64, + ) -> i32; /// The whole `hc_post` re-expansion. Same return contract as above. #[allow(clippy::too_many_arguments)] pub(crate) fn hc_post_fused( diff --git a/mistralrs-core/src/cuda/hc_fused.cu b/mistralrs-core/src/cuda/hc_fused.cu index 2ff941e57..29f0c629c 100644 --- a/mistralrs-core/src/cuda/hc_fused.cu +++ b/mistralrs-core/src/cuda/hc_fused.cu @@ -495,6 +495,266 @@ __global__ void hc_post_fused_kernel( } } +// --------------------------------------------------------------------------- +// THE F32 SEAM (3 kernels below) +// --------------------------------------------------------------------------- +// +// Everything above collapses chains that were *already* F32 end to end. The +// three kernels below close the other half of the problem: expressions that +// are BF16 on both sides but were spelled as "widen -> a few elementwise ops +// in F32 -> narrow". Every entry and exit paid a `cast_*` launch and every +// intermediate was a materialised F32 tensor. Measured on this H200 at +// b424f5cd9 (agent/datamove-fuse), per decode forward: +// +// cast_bf16_f32 356.7 | cast_f32_bf16 221.1 | bminimum_f32 172.0 +// bmaximum_f32 86.0 | bmul_f32 94.4 | fast_sum_f32 98.8 +// +// `bmaximum_f32 = 86.0` and `bminimum_f32 = 172.0` are exact, not approximate: +// 86 = 2 x 43 layers is the number of `swiglu_clamp` calls per forward (one +// routed-expert stage + one shared-expert stage per layer), and each contains +// one `maximum` and two `minimum`. That fingerprint is what identified the +// seam. +// +// A SECOND, INVISIBLE COST. `Tensor::minimum(f64)` / `maximum(f64)` build the +// scalar operand as `Tensor::new(v, Cpu)?.to_dtype()?.to_device(cuda)?` on +// every call (candle tensor.rs `binary_op_scalar!`). That is a 4-byte +// host-to-device memcpy per clamp bound -- 3 per `swiglu_clamp`, 258 per +// forward -- which never appears in a kernel histogram at all. +// +// BIT-IDENTITY. Same contract as the rest of this file: transcription, not +// re-derivation. The op-by-op sources are quoted at each kernel. `ARC_SEAM_AB=1` +// re-runs the candle chain next to each kernel and compares raw bits, and it +// compares them at **F32, before the narrowing** -- a BF16-only comparison has +// 8 mantissa bits and would swallow exactly the reassociation errors that +// matter. `ARC_HC_AB_POISON=1` perturbs one element by 1 ULP and must make +// every comparison fail; if it does not, the comparison is broken, not the +// kernel. + +// V4 clamped SwiGLU: 8 launches + 3 H2D memcpys -> 1 kernel. +// +// Replaces verbatim, for the ROUTED experts (moe/experts.rs `swiglu` -> +// `swiglu_clamp`): +// let gate = gate.to_dtype(F32)?.minimum(limit)?; // cast + bminimum +// let up = up.to_dtype(F32)?.clamp(-limit, limit)?; // cast + bmaximum + bminimum +// let act = gate.apply(&Silu)?; // usilu_f32 +// up.mul(&act)?.to_dtype(out_dtype) // bmul + cast +// +// and for the SHARED expert (layers.rs `Mlp::gated_act`), whose middle two +// launches are instead one `fused_glu_f32` (mistralrs-quant ops.cu:845-890): +// T activated = (T)glu_silu(a); output[i] = activated * b[i]; +// +// The two spellings are the SAME VALUE, which is why one kernel serves both: +// * candle `usilu_f32` is `silu_fwd(x) = x / (static_cast(1) + expg(-x))` +// (candle-kernels unary.cu:59-61, cuda_utils.cuh:148 `expg(float)=expf`); +// * `glu_silu(x)` is `x / (1.0f + expf(-x))` -- character for character the +// same expression; +// * the products differ only in operand order (`up * act` vs `act * up`), +// and IEEE multiplication is commutative bit-for-bit. +// +// `clamp(lo, hi)` is `maximum(lo)?.minimum(hi)` (candle tensor.rs:1176-1178), +// and candle's `maxg`/`ming` for float are `fmaxf`/`fminf` +// (cuda_utils.cuh:142-145) -- reproduced, not re-derived. The clamp bounds +// arrive as f64 in Rust and are narrowed to f32 by candle's `to_dtype` before +// the compare; `limit` is passed in already narrowed, so `-limit` here is the +// same f32 the eager path compares against. +// +// TIN/TOUT are separate so `ARC_SEAM_AB=1` can ask for the F32 product that +// the eager chain holds *before* its final cast, and compare there. +template +__global__ void arc_seam_swiglu_clamp_kernel( + const TIN *__restrict__ gate, // [total] model dtype + const TIN *__restrict__ up, // [total] model dtype + TOUT *__restrict__ out, // [total] + float limit, + long long total +) { + const long long stride = (long long)blockDim.x * (long long)gridDim.x; + const float neg_limit = -limit; + for (long long i = (long long)blockIdx.x * (long long)blockDim.x + threadIdx.x; + i < total; i += stride) { + // BF16->F32 is an exact widening: these are the bits candle's + // `to_dtype(F32)` would have written to a temporary. + float g = hc_widen(gate[i]); + float u = hc_widen(up[i]); + g = fminf(g, limit); // gate is clamped ONLY from above + u = fmaxf(u, neg_limit); // up is clamped on both sides, max first + u = fminf(u, limit); + // silu_fwd: one add, one accurate-libdevice expf, one div.rn.f32. + // There is no mul+add pair here, so --fmad=false has nothing to + // forbid and the result does not depend on the contraction setting. + const float den = __fadd_rn(1.0f, expf(-g)); + const float act = __fdiv_rn(g, den); + out[i] = hc_narrow(__fmul_rn(act, u)); + } +} + +// The clamp HALF of the same stage: 5 launches + 3 H2D memcpys -> 1 kernel, +// two outputs, both F32. +// +// Why a second, weaker kernel exists. The SHARED expert's activation and +// multiply are `mistralrs_quant::fused_glu`, and mistralrs-quant/build.rs:281 +// compiles that translation unit with `--use_fast_math`. Its `expf` is +// therefore `__expf` (ex2.approx) and its `/` is div.approx, not the +// `__nv_expf` + div.rn that candle-kernels' `usilu_f32` uses -- and fast math +// also implies -ftz=true, which no intrinsic in a non-fast-math file can +// reproduce per-op. So `arc_seam_swiglu_clamp_kernel` above is bit-identical to +// the ROUTED chain and would NOT be bit-identical to the shared one. Rather +// than fuse the shared site on a tolerance argument, only its provable half is +// fused: the two widening casts, the `maximum` and the two `minimum`s, plus the +// three host-to-device copies of the bounds. `fused_glu` then runs unchanged on +// the F32 tensors it runs on today. +// +// `ARC_SEAM_AB=1` still evaluates the full fusion at the shared site and +// reports whether the fast-math and IEEE spellings happen to agree; that +// measurement decides whether the remaining 4 launches per shared call can ever +// be taken, and it is a fact in the log rather than an assumption either way. +template +__global__ void arc_seam_swiglu_clamp_split_kernel( + const TIN *__restrict__ gate, // [total] model dtype + const TIN *__restrict__ up, // [total] model dtype + float *__restrict__ gate_out, // [total] F32, min(g, limit) + float *__restrict__ up_out, // [total] F32, clamp(u, -limit, limit) + float limit, + long long total +) { + const long long stride = (long long)blockDim.x * (long long)gridDim.x; + const float neg_limit = -limit; + for (long long i = (long long)blockIdx.x * (long long)blockDim.x + threadIdx.x; + i < total; i += stride) { + const float g = hc_widen(gate[i]); + float u = hc_widen(up[i]); + u = fmaxf(u, neg_limit); + u = fminf(u, limit); + gate_out[i] = fminf(g, limit); + up_out[i] = u; + } +} + +// MoE expert combine: 4 launches -> 1. +// +// Replaces verbatim (moe/experts.rs `forward_fast`, the `experts.weighted_sum` +// span): +// ys.to_dtype(F32)? // cast_bf16_f32 +// .broadcast_mul(&topk_weights.unsqueeze(-1)?)? // bmul_f32 +// .sum(D::Minus2)? // fast_sum_f32 +// .to_dtype(original_dtype) // cast_f32_bf16 +// +// REDUCTION ORDER. `sum(D::Minus2)` on `[n, K, h]` gives dst_el = n*h and +// el_to_sum_per_block = K, so candle's `FastReduce` picks +// block_dim = min(1024, K).next_power_of_two() = P (cuda_backend/mod.rs:370). +// For V4's K = 6 that is P = 8: threads 0..5 each take one element, threads +// 6 and 7 take none and keep their zero-initialised accumulator, and the +// pairwise tree then runs over all EIGHT slots. The sum is therefore +// ((0+a0)+(0+a4)) + ((0+a2)+0) + [ ((0+a1)+(0+a5)) + ((0+a3)+0) ] +// which is neither the sequential sum nor the K=6 pairwise tree. The identity +// padding is part of the answer and is replayed literally below, as is the +// `shr[tid] = 0; shr[tid] += v` (it canonicalises -0.0f to +0.0f). +template +__global__ void arc_seam_moe_weighted_sum_kernel( + const TIN *__restrict__ ys, // [n, K, h] model dtype + const float *__restrict__ w, // [n, K] F32 routing weights + TOUT *__restrict__ out, // [n, h] + int h, + long long total // n * h +) { + const long long stride = (long long)blockDim.x * (long long)gridDim.x; + for (long long i = (long long)blockIdx.x * (long long)blockDim.x + threadIdx.x; + i < total; i += stride) { + const long long row = i / (long long)h; + const int col = (int)(i - row * (long long)h); + const float *wrow = w + row * K; + const TIN *yrow = ys + row * (long long)K * (long long)h + col; + + float acc[P]; +#pragma unroll + for (int t = 0; t < P; ++t) { + if (t < K) { + // `broadcast_mul` is binary.cu `x * y` with x = ys, y = w. + acc[t] = __fadd_rn( + 0.0f, + __fmul_rn(hc_widen(yrow[(long long)t * (long long)h]), wrow[t])); + } else { + acc[t] = 0.0f; // identity-padded lane: shr[tid] = 0, never added to + } + } +#pragma unroll + for (int s = P >> 1; s > 0; s >>= 1) { +#pragma unroll + for (int t = 0; t < P; ++t) { + if (t < s) { + acc[t] = __fadd_rn(acc[t], acc[t + s]); + } + } + } + out[i] = hc_narrow(acc[0]); + } +} + +// MoE gate weight renormalise + scale: 4 launches -> 1. +// +// Replaces verbatim (deepseek4.rs `MoeGate::forward`, spans `gate.renormalize`): +// let denominator = (topk_weight.sum_keepdim(D::Minus1)? + 1e-20)?; // fast_sum + affine +// topk_weight = topk_weight.broadcast_div(&denominator)?; // bdiv_f32 +// topk_weight = (topk_weight * routed_scaling_factor)?; // affine +// +// `sum_keepdim(D::Minus1)` on `[n, K]` gives el_to_sum_per_block = K, so the +// same P = next_pow2(K) padded tree as above applies. +// +// `Tensor + f64` and `Tensor * f64` both lower to `affine(mul, add)` +// (candle tensor.rs `bin_trait!`: Add -> affine(1., v), Mul -> affine(v, 0.)), +// whose kernel body is `x * mul + add` (affine.cu:18) compiled by candle with +// the DEFAULT -fmad=true, i.e. a single fmaf. Written as an explicit `fmaf` +// here: an explicit fmaf is unaffected by this file's --fmad=false, which only +// forbids *contraction* of a separate mul and add. `fmaf(x, scale, 0.0f)` is +// deliberately not `x * scale` -- they differ in the sign of zero. +template +__global__ void arc_seam_gate_renorm_kernel( + const float *__restrict__ w, // [n, K] + float *__restrict__ out, // [n, K] + float eps, // the 1e-20 added to the denominator + float scale, // routed_scaling_factor + int do_renorm, // cfg.norm_topk_prob && scoring in {sigmoid, sqrtsoftplus} + long long n +) { + const long long stride = (long long)blockDim.x * (long long)gridDim.x; + for (long long row = (long long)blockIdx.x * (long long)blockDim.x + threadIdx.x; + row < n; row += stride) { + const float *wrow = w + row * K; + float *orow = out + row * K; + + float v[K]; +#pragma unroll + for (int t = 0; t < K; ++t) { + v[t] = wrow[t]; + } + + float den = 0.0f; + if (do_renorm) { + float acc[P]; +#pragma unroll + for (int t = 0; t < P; ++t) { + acc[t] = (t < K) ? __fadd_rn(0.0f, v[t]) : 0.0f; + } +#pragma unroll + for (int s = P >> 1; s > 0; s >>= 1) { +#pragma unroll + for (int t = 0; t < P; ++t) { + if (t < s) { + acc[t] = __fadd_rn(acc[t], acc[t + s]); + } + } + } + den = fmaf(acc[0], 1.0f, eps); // affine(1., 1e-20) + } + +#pragma unroll + for (int t = 0; t < K; ++t) { + const float x = do_renorm ? __fdiv_rn(v[t], den) : v[t]; + orow[t] = fmaf(x, scale, 0.0f); // affine(routed_scaling_factor, 0.) + } + } +} + // 256 threads, grid capped so the grid-stride loop stays valid for any size. inline void hc_launch_dims(long long total, unsigned &blocks, int &block) { block = 256; @@ -598,4 +858,154 @@ int hc_post_fused( #undef HC_PF } +// --------------------------------------------------------------------------- +// F32-seam entry points. +// +// Deliberately fresh symbol names. A stale object file left in the archive by +// an incremental build has already, once, satisfied a link with none of the +// new code in it -- byte-identical output, green build, silent no-op. A name +// that has never existed before turns that failure into a link error. +// --------------------------------------------------------------------------- + +// `out = narrow( silu(min(gate, limit)) * clamp(up, -limit, limit) )`. +// `in_dtype` is the dtype of gate/up; `out_dtype` is usually the same, and is +// F32 only when the A/B harness asks for the pre-narrowing product. +int arc_seam_swiglu_clamp( + const void *gate, + const void *up, + void *out, + float limit, + long long total, + int in_dtype, + int out_dtype, + long long stream +) { + unsigned blocks; + int block; + hc_launch_dims(total, blocks, block); + cudaStream_t s = (cudaStream_t)stream; +#define ARC_SWG(TIN, TOUT) \ + arc_seam_swiglu_clamp_kernel<<>>( \ + (const TIN *)gate, (const TIN *)up, (TOUT *)out, limit, total) + if (in_dtype == HC_DTYPE_BF16 && out_dtype == HC_DTYPE_BF16) { + ARC_SWG(__nv_bfloat16, __nv_bfloat16); + return 0; + } + if (in_dtype == HC_DTYPE_BF16 && out_dtype == HC_DTYPE_F32) { + ARC_SWG(__nv_bfloat16, float); + return 0; + } + if (in_dtype == HC_DTYPE_F32 && out_dtype == HC_DTYPE_F32) { + ARC_SWG(float, float); + return 0; + } + if (in_dtype == HC_DTYPE_F32 && out_dtype == HC_DTYPE_BF16) { + ARC_SWG(float, __nv_bfloat16); + return 0; + } + return 1; +#undef ARC_SWG +} + +// `gate_out = min(gate, limit)`, `up_out = clamp(up, -limit, limit)`, both F32. +int arc_seam_swiglu_clamp_split( + const void *gate, + const void *up, + void *gate_out, + void *up_out, + float limit, + long long total, + int in_dtype, + long long stream +) { + unsigned blocks; + int block; + hc_launch_dims(total, blocks, block); + cudaStream_t s = (cudaStream_t)stream; +#define ARC_SWG_SPLIT(TIN) \ + arc_seam_swiglu_clamp_split_kernel<<>>( \ + (const TIN *)gate, (const TIN *)up, (float *)gate_out, (float *)up_out, limit, total) + if (in_dtype == HC_DTYPE_BF16) { + ARC_SWG_SPLIT(__nv_bfloat16); + return 0; + } + if (in_dtype == HC_DTYPE_F32) { + ARC_SWG_SPLIT(float); + return 0; + } + return 1; +#undef ARC_SWG_SPLIT +} + +// `out[n, h] = narrow( tree_sum_K( ys[n, j, h] * w[n, j] ) )`. +int arc_seam_moe_weighted_sum( + const void *ys, + const void *w, + void *out, + int k, + int h, + long long total, + int in_dtype, + int out_dtype, + long long stream +) { + unsigned blocks; + int block; + hc_launch_dims(total, blocks, block); + cudaStream_t s = (cudaStream_t)stream; +#define ARC_WS(TIN, TOUT, KV, PV) \ + arc_seam_moe_weighted_sum_kernel<<>>( \ + (const TIN *)ys, (const float *)w, (TOUT *)out, h, total) +#define ARC_WS_K(TIN, TOUT) \ + switch (k) { \ + case 1: ARC_WS(TIN, TOUT, 1, 1); return 0; \ + case 2: ARC_WS(TIN, TOUT, 2, 2); return 0; \ + case 4: ARC_WS(TIN, TOUT, 4, 4); return 0; \ + case 6: ARC_WS(TIN, TOUT, 6, 8); return 0; \ + case 8: ARC_WS(TIN, TOUT, 8, 8); return 0; \ + default: return 1; \ + } + if (in_dtype == HC_DTYPE_BF16 && out_dtype == HC_DTYPE_BF16) { + ARC_WS_K(__nv_bfloat16, __nv_bfloat16) + } else if (in_dtype == HC_DTYPE_BF16 && out_dtype == HC_DTYPE_F32) { + ARC_WS_K(__nv_bfloat16, float) + } else if (in_dtype == HC_DTYPE_F32 && out_dtype == HC_DTYPE_F32) { + ARC_WS_K(float, float) + } else if (in_dtype == HC_DTYPE_F32 && out_dtype == HC_DTYPE_BF16) { + ARC_WS_K(float, __nv_bfloat16) + } + return 1; +#undef ARC_WS_K +#undef ARC_WS +} + +// `out[n, :] = affine( w[n, :] / (tree_sum_K(w[n, :]) + eps), scale )`. +int arc_seam_gate_renorm( + const void *w, + void *out, + float eps, + float scale, + int do_renorm, + int k, + long long n, + long long stream +) { + unsigned blocks; + int block; + hc_launch_dims(n, blocks, block); + cudaStream_t s = (cudaStream_t)stream; +#define ARC_GR(KV, PV) \ + arc_seam_gate_renorm_kernel<<>>( \ + (const float *)w, (float *)out, eps, scale, do_renorm, n) + switch (k) { + case 1: ARC_GR(1, 1); return 0; + case 2: ARC_GR(2, 2); return 0; + case 4: ARC_GR(4, 4); return 0; + case 6: ARC_GR(6, 8); return 0; + case 8: ARC_GR(8, 8); return 0; + default: return 1; + } +#undef ARC_GR +} + } // extern "C" diff --git a/mistralrs-core/src/cuda/hc_fused.rs b/mistralrs-core/src/cuda/hc_fused.rs index d893340c8..928994ed9 100644 --- a/mistralrs-core/src/cuda/hc_fused.rs +++ b/mistralrs-core/src/cuda/hc_fused.rs @@ -279,6 +279,185 @@ mod cuda_impl { } } + /// Fused V4 clamped SwiGLU: `narrow(silu(min(gate, limit)) * clamp(up, -limit, limit))`. + /// + /// Replaces the 8-launch (routed) / 7-launch (shared) chain in + /// `moe::experts::swiglu_clamp` + the activation stage that follows it, + /// plus the three 4-byte host-to-device copies candle makes for the clamp + /// bounds. + /// + /// - `gate`, `up`: same shape, same dtype (F32 or BF16), contiguous, CUDA + /// - `out_dtype`: normally `gate.dtype()`; F32 is used by the A/B harness + /// to obtain the product *before* the narrowing, where a bitwise + /// comparison is not swallowed by BF16's 8 mantissa bits. + pub fn swiglu_clamp_fused_cuda( + gate: &Tensor, + up: &Tensor, + limit: f32, + out_dtype: DType, + ) -> Result { + if gate.dims() != up.dims() { + candle::bail!( + "hc_fused: swiglu gate {:?} vs up {:?}", + gate.dims(), + up.dims() + ); + } + if gate.dtype() != up.dtype() { + candle::bail!( + "hc_fused: swiglu gate is {:?} but up is {:?}", + gate.dtype(), + up.dtype() + ); + } + let total = gate.elem_count(); + if total == 0 { + candle::bail!("hc_fused: empty swiglu"); + } + let in_tag = dtype_tag(gate.dtype())?; + let out_tag = dtype_tag(out_dtype)?; + let gate_ptr = any_ptr(gate, "swiglu gate")?; + let up_ptr = any_ptr(up, "swiglu up")?; + let dev = gate.device().as_cuda_device()?; + let stream = dev.cuda_stream().cu_stream() as i64; + let shape = gate.shape().clone(); + + #[allow(clippy::cast_possible_truncation)] + alloc_out(&dev, out_dtype, total, shape, |out| unsafe { + crate::cuda::ffi::arc_seam_swiglu_clamp( + gate_ptr, up_ptr, out, limit, total as i64, in_tag, out_tag, stream, + ) + }) + } + + /// The clamp half of [`swiglu_clamp_fused_cuda`], returning + /// `(min(gate, limit), clamp(up, -limit, limit))` in F32. + /// + /// Used at the shared-expert site, whose activation is `fused_glu` compiled + /// with `--use_fast_math` (mistralrs-quant/build.rs:281) and therefore not + /// reproducible from this no-fast-math translation unit. See the kernel's + /// comment in `hc_fused.cu` for why the weaker fusion is the honest one + /// there. + pub fn swiglu_clamp_split_cuda( + gate: &Tensor, + up: &Tensor, + limit: f32, + ) -> Result<(Tensor, Tensor)> { + use candle_core::cuda_backend::cudarc::driver::DevicePtr; + + if gate.dims() != up.dims() { + candle::bail!( + "hc_fused: swiglu-split gate {:?} vs up {:?}", + gate.dims(), + up.dims() + ); + } + if gate.dtype() != up.dtype() { + candle::bail!("hc_fused: swiglu-split gate/up dtypes differ"); + } + let total = gate.elem_count(); + if total == 0 { + candle::bail!("hc_fused: empty swiglu-split"); + } + let in_tag = dtype_tag(gate.dtype())?; + let gate_ptr = any_ptr(gate, "swiglu gate")?; + let up_ptr = any_ptr(up, "swiglu up")?; + let dev = gate.device().as_cuda_device()?; + let stream = dev.cuda_stream().cu_stream() as i64; + let shape = gate.shape().clone(); + + let g_buf = unsafe { dev.alloc::(total) }?; + let u_buf = unsafe { dev.alloc::(total) }?; + #[allow(clippy::cast_possible_truncation)] + let rc = unsafe { + crate::cuda::ffi::arc_seam_swiglu_clamp_split( + gate_ptr, + up_ptr, + g_buf.device_ptr(g_buf.stream()).0 as *mut std::ffi::c_void, + u_buf.device_ptr(u_buf.stream()).0 as *mut std::ffi::c_void, + limit, + total as i64, + in_tag, + stream, + ) + }; + if rc != 0 { + candle::bail!("hc_fused: swiglu-split dispatch refused (rc={rc})"); + } + let g_st = candle::CudaStorage::wrap_cuda_slice(g_buf, dev.clone()); + let u_st = candle::CudaStorage::wrap_cuda_slice(u_buf, dev.clone()); + Ok(( + Tensor::from((candle::Storage::Cuda(g_st), shape.clone())), + Tensor::from((candle::Storage::Cuda(u_st), shape)), + )) + } + + /// Fused MoE expert combine: `out[n, h] = narrow(sum_j ys[n, j, h] * w[n, j])`, + /// replacing `to_dtype + broadcast_mul + sum(-2) + to_dtype` (4 launches). + /// + /// - `ys` `[n, k, h]` model dtype, contiguous + /// - `w` `[n, k]` F32, contiguous + pub fn moe_weighted_sum_cuda(ys: &Tensor, w: &Tensor, out_dtype: DType) -> Result { + let (n, k, h) = ys.dims3()?; + let (wn, wk) = w.dims2()?; + if wn != n || wk != k { + candle::bail!("hc_fused: weights are [{wn}, {wk}] but ys is [{n}, {k}, {h}]"); + } + if n == 0 || h == 0 { + candle::bail!("hc_fused: empty weighted sum"); + } + let in_tag = dtype_tag(ys.dtype())?; + let out_tag = dtype_tag(out_dtype)?; + let ys_ptr = any_ptr(ys, "ys")?; + let w_ptr = f32_ptr(w, "topk_weights")?; + let dev = ys.device().as_cuda_device()?; + let stream = dev.cuda_stream().cu_stream() as i64; + let total = (n * h) as i64; + + #[allow(clippy::cast_possible_truncation)] + alloc_out(&dev, out_dtype, n * h, (n, h), |out| unsafe { + crate::cuda::ffi::arc_seam_moe_weighted_sum( + ys_ptr, w_ptr, out, k as i32, h as i32, total, in_tag, out_tag, stream, + ) + }) + } + + /// Fused MoE gate renormalise + scale, replacing + /// `sum_keepdim + affine(+eps) + broadcast_div + affine(*scale)` + /// (4 launches). + /// + /// `w` is `[n, k]` F32 contiguous. `do_renorm == false` applies only the + /// scale, which is what the config-gated eager path does when + /// `norm_topk_prob` is off. + pub fn gate_renorm_cuda( + w: &Tensor, + eps: f32, + scale: f32, + do_renorm: bool, + ) -> Result { + let (n, k) = w.dims2()?; + if n == 0 || k == 0 { + candle::bail!("hc_fused: empty gate renorm"); + } + let w_ptr = f32_ptr(w, "topk_weight")?; + let dev = w.device().as_cuda_device()?; + let stream = dev.cuda_stream().cu_stream() as i64; + + #[allow(clippy::cast_possible_truncation)] + alloc_out(&dev, DType::F32, n * k, (n, k), |out| unsafe { + crate::cuda::ffi::arc_seam_gate_renorm( + w_ptr, + out, + eps, + scale, + i32::from(do_renorm), + k as i32, + n as i64, + stream, + ) + }) + } + /// Fused `hc_pre` tail: `y = sum_i pre[i] * x_f32[i, :]`, narrowed to /// `out_dtype`. Replaces `broadcast_mul + sum(1) + to_dtype` (3 launches). /// @@ -360,7 +539,48 @@ mod cuda_impl { } #[cfg(feature = "cuda")] -pub use cuda_impl::{hc_post_fused_cuda, hc_pre_fused_cuda, hc_y_combine_cuda, sqrt_softplus_cuda}; +pub use cuda_impl::{ + gate_renorm_cuda, hc_post_fused_cuda, hc_pre_fused_cuda, hc_y_combine_cuda, + moe_weighted_sum_cuda, sqrt_softplus_cuda, swiglu_clamp_fused_cuda, swiglu_clamp_split_cuda, +}; + +#[cfg(not(feature = "cuda"))] +pub fn swiglu_clamp_split_cuda( + _gate: &candle_core::Tensor, + _up: &candle_core::Tensor, + _limit: f32, +) -> candle_core::Result<(candle_core::Tensor, candle_core::Tensor)> { + candle_core::bail!("swiglu_clamp_split_cuda requires the cuda feature") +} + +#[cfg(not(feature = "cuda"))] +pub fn swiglu_clamp_fused_cuda( + _gate: &candle_core::Tensor, + _up: &candle_core::Tensor, + _limit: f32, + _out_dtype: candle_core::DType, +) -> candle_core::Result { + candle_core::bail!("swiglu_clamp_fused_cuda requires the cuda feature") +} + +#[cfg(not(feature = "cuda"))] +pub fn moe_weighted_sum_cuda( + _ys: &candle_core::Tensor, + _w: &candle_core::Tensor, + _out_dtype: candle_core::DType, +) -> candle_core::Result { + candle_core::bail!("moe_weighted_sum_cuda requires the cuda feature") +} + +#[cfg(not(feature = "cuda"))] +pub fn gate_renorm_cuda( + _w: &candle_core::Tensor, + _eps: f32, + _scale: f32, + _do_renorm: bool, +) -> candle_core::Result { + candle_core::bail!("gate_renorm_cuda requires the cuda feature") +} #[cfg(not(feature = "cuda"))] pub fn hc_y_combine_cuda( @@ -429,6 +649,119 @@ pub fn usable_model_dtype(t: &candle_core::Tensor) -> bool { && fused_enabled() } +// --------------------------------------------------------------------------- +// The F32 seam — gates and engagement accounting +// --------------------------------------------------------------------------- +// +// One master switch so the whole change can be A/B'd from a single binary +// (`ARC_F32SEAM=0` restores every eager chain), plus one switch per site so a +// site that turns out to cost bit-identity can be dropped without a rebuild. +// +// ENGAGEMENT IS REPORTED, NOT ASSUMED. A timing win from a fast path that never +// ran is this codebase's most-repeated failure. Every site logs the first time +// it engages AND the first time it declines, with the reason, and prints running +// totals; "no DECLINED line" is then a fact in the log rather than an inference. + +/// `ARC_F32SEAM=0` restores every eager chain in the seam. Anything else keeps +/// them fused. +pub fn seam_enabled() -> bool { + static ON: std::sync::OnceLock = std::sync::OnceLock::new(); + *ON.get_or_init(|| !matches!(std::env::var("ARC_F32SEAM").as_deref(), Ok("0"))) +} + +/// Per-site opt-out, e.g. `ARC_SEAM_SWIGLU=0`. Only consulted when the master +/// switch is on. +fn seam_site_enabled(var: &'static str) -> bool { + // Sites are few and fixed; a tiny linear scan of cached values beats a + // per-call `std::env::var` (which takes a global lock). + static CACHE: std::sync::OnceLock> = std::sync::OnceLock::new(); + let cache = CACHE.get_or_init(|| { + ["ARC_SEAM_SWIGLU", "ARC_SEAM_WSUM", "ARC_SEAM_GATE"] + .into_iter() + .map(|v| (v, !matches!(std::env::var(v).as_deref(), Ok("0")))) + .collect() + }); + cache + .iter() + .find(|(name, _)| *name == var) + .map(|(_, on)| *on) + .unwrap_or(true) +} + +/// Whether a seam site should fuse: master on, this site on, CUDA build. +pub fn seam_on(var: &'static str) -> bool { + cfg!(feature = "cuda") && seam_enabled() && seam_site_enabled(var) +} + +/// `ARC_SEAM_AB=1` recomputes the eager chain beside every seam kernel and +/// compares raw bits — at **F32**, before any narrowing, because a BF16 +/// comparison has 8 mantissa bits and would swallow the reassociation errors +/// this is meant to catch. Shares `ARC_HC_AB_POISON=1` as its negative control. +pub fn seam_ab_enabled() -> bool { + static ON: std::sync::OnceLock = std::sync::OnceLock::new(); + *ON.get_or_init(|| matches!(std::env::var("ARC_SEAM_AB").as_deref(), Ok("1"))) +} + +static SEAM_ENGAGED: [AtomicU64; 4] = [ + AtomicU64::new(0), + AtomicU64::new(0), + AtomicU64::new(0), + AtomicU64::new(0), +]; +static SEAM_DECLINED: [AtomicU64; 4] = [ + AtomicU64::new(0), + AtomicU64::new(0), + AtomicU64::new(0), + AtomicU64::new(0), +]; + +/// Site indices into the counters above. +pub const SEAM_SWIGLU: usize = 0; +pub const SEAM_WSUM: usize = 1; +pub const SEAM_GATE: usize = 2; +pub const SEAM_SWIGLU_SPLIT: usize = 3; + +const SEAM_NAMES: [&str; 4] = [ + "swiglu_clamp", + "moe_weighted_sum", + "gate_renorm", + "swiglu_clamp_split", +]; + +/// Record that a seam site fused. Logs the first one and then periodic totals. +pub fn seam_engaged(site: usize) { + let prev = SEAM_ENGAGED[site].fetch_add(1, Ordering::Relaxed); + if prev == 0 { + tracing::warn!("arc-f32-seam ENGAGED site={}", SEAM_NAMES[site]); + } else if (prev + 1) % 10_000 == 0 { + tracing::warn!( + "arc-f32-seam site={} engaged={} declined={}", + SEAM_NAMES[site], + prev + 1, + SEAM_DECLINED[site].load(Ordering::Relaxed) + ); + } +} + +/// Record that a seam site fell back, and why. Logs the first one per site. +pub fn seam_declined(site: usize, reason: &str) { + let prev = SEAM_DECLINED[site].fetch_add(1, Ordering::Relaxed); + if prev == 0 { + tracing::warn!( + "arc-f32-seam DECLINED site={} reason={reason}", + SEAM_NAMES[site] + ); + } +} + +/// `(engaged, declined)` totals for a site — used by tests and by the report. +pub fn seam_counts(site: usize) -> (u64, u64) { + ( + SEAM_ENGAGED[site].load(Ordering::Relaxed), + SEAM_DECLINED[site].load(Ordering::Relaxed), + ) +} + // --------------------------------------------------------------------------- // The on-GPU bitwise A/B — `ARC_HC_AB=1` // --------------------------------------------------------------------------- @@ -889,6 +1222,134 @@ pub(crate) mod reference { } } +#[cfg(test)] +mod seam_tests { + use super::*; + + /// Scalar replica of `arc_seam_moe_weighted_sum_kernel` / + /// `arc_seam_gate_renorm_kernel`'s reduction: the identity-padded pairwise + /// tree over `P = next_pow2(K)` slots. + fn seam_padded_tree(vals: &[f32]) -> f32 { + let p = vals.len().next_power_of_two(); + let mut acc = vec![0.0f32; p]; + for (t, v) in vals.iter().enumerate() { + acc[t] = 0.0f32 + v; // the `shr[tid] = 0; shr[tid] += v` of fast_sum + } + let mut s = p / 2; + while s > 0 { + for t in 0..s { + acc[t] += acc[t + s]; + } + s >>= 1; + } + acc[0] + } + + /// The seam kernels must reduce in candle's `fast_sum` order, which for + /// V4's `top_k = 6` is a tree over EIGHT slots, not six. This is the whole + /// bit-identity risk of the two reduction kernels, and it is checkable + /// without a GPU: `sinkhorn::reference::candle_tree_sum` is an independent + /// transcription of the same candle kernel, written for a different fusion. + #[test] + fn seam_reduction_replays_candle_fast_sum_including_the_padding_lanes() { + let mut seed = 0x5eedu64; + let mut next = || { + seed = seed + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + (((seed >> 40) as f32) / ((1u32 << 24) as f32)) * 2.0 - 1.0 + }; + for _ in 0..2000 { + for k in [1usize, 2, 4, 6, 8] { + let vals: Vec = (0..k).map(|_| next()).collect(); + assert_eq!( + seam_padded_tree(&vals).to_bits(), + crate::cuda::sinkhorn::reference::candle_tree_sum(&vals).to_bits(), + "k={k} vals={vals:?}" + ); + } + } + } + + /// For `k = 6` the padded tree is NOT the sequential sum, and this pins the + /// exact association the kernel must reproduce. A kernel that summed + /// `a0..a5` left to right would pass a tolerance test and fail this one. + #[test] + fn k6_tree_is_not_the_sequential_sum() { + // The tree pairs a0 with a4, so the two large opposite-sign terms + // cancel exactly there while a left-to-right sum absorbs the small + // terms into 1e8 and loses them. 1.0 vs 4.0 — not a last-ulp + // difference, a catastrophic one. + let v = [1e8f32, 1.0, 1.0, 1.0, -1e8, 1.0]; + let seq = v.iter().fold(0.0f32, |a, b| a + b); + let tree = seam_padded_tree(&v); + assert_ne!( + seq.to_bits(), + tree.to_bits(), + "if these ever agree this test has stopped discriminating" + ); + // ((0+a0)+(0+a4)) + ((0+a2)+0) + ((0+a1)+(0+a5)) + ((0+a3)+0) + let lhs = (v[0] + v[4]) + v[2]; + let rhs = (v[1] + v[5]) + v[3]; + assert_eq!(tree.to_bits(), (lhs + rhs).to_bits()); + } + + /// `fmaf(x, scale, 0.0)` is what candle's `affine(mul, 0.)` compiles to and + /// is deliberately not `x * scale`: they disagree on the sign of zero, and + /// the gate renorm kernel writes the result of exactly this expression. + #[test] + fn affine_is_an_fma_not_a_multiply() { + let x = -0.0f32; + assert_ne!( + (x * 2.0f32).to_bits(), + x.mul_add(2.0, 0.0).to_bits(), + "the two spellings must still differ, or this guard proves nothing" + ); + assert_eq!(x.mul_add(2.0, 0.0).to_bits(), 0.0f32.to_bits()); + } + + /// Engagement accounting must actually count, in both directions — + /// otherwise "no DECLINED line in the log" is not evidence of anything. + #[test] + fn seam_engagement_counters_move() { + let before = seam_counts(SEAM_GATE); + seam_engaged(SEAM_GATE); + seam_declined(SEAM_GATE, "unit test"); + let after = seam_counts(SEAM_GATE); + assert_eq!(after.0, before.0 + 1); + assert_eq!(after.1, before.1 + 1); + } + + /// The seam kernels must stay in the no-fast-math translation unit, and the + /// shared-expert site must stay on the weaker fusion for as long as + /// `fused_glu` is compiled with `--use_fast_math`. If that flag ever leaves + /// mistralrs-quant/build.rs, the full fusion becomes legal there and this + /// test is the reminder. + #[test] + fn seam_kernels_are_where_the_bit_identity_argument_says_they_are() { + let cu = include_str!("hc_fused.cu"); + for sym in [ + "arc_seam_swiglu_clamp_kernel", + "arc_seam_swiglu_clamp_split_kernel", + "arc_seam_moe_weighted_sum_kernel", + "arc_seam_gate_renorm_kernel", + ] { + assert!(cu.contains(sym), "{sym} left hc_fused.cu"); + } + let build = include_str!("../../build.rs"); + assert!( + build.contains("src/cuda/hc_fused.cu"), + "hc_fused.cu must stay in the dedicated no-fast-math builder" + ); + let quant_build = include_str!("../../../mistralrs-quant/build.rs"); + assert!( + quant_build.contains("--use_fast_math"), + "if mistralrs-quant stopped using fast math, fused_glu became \ + reproducible and the shared expert can take the full fusion" + ); + } +} + #[cfg(test)] mod tests { use super::reference::*; diff --git a/mistralrs-core/src/layers.rs b/mistralrs-core/src/layers.rs index 39eefbad2..3a07b09fb 100644 --- a/mistralrs-core/src/layers.rs +++ b/mistralrs-core/src/layers.rs @@ -3242,7 +3242,15 @@ impl Mlp { return crate::ops::mul_and_act(gate, up, self.act); }; let out_dtype = gate.dtype(); - let (gate, up) = crate::moe::swiglu_clamp(gate, up, limit)?; + // The shared-expert half of the 86 clamp calls per decode forward. Only + // the CLAMP is fused here, not the activation: `mul_and_act` below + // lands in `mistralrs_quant::fused_glu`, whose translation unit is + // built with `--use_fast_math`, so a no-fast-math kernel cannot + // reproduce its bits. See `crate::moe::fused_swiglu_clamp_split`. + let (gate, up) = match crate::moe::fused_swiglu_clamp_split(gate, up, limit, self.act)? { + Some(pair) => pair, + None => crate::moe::swiglu_clamp(gate, up, limit)?, + }; crate::ops::mul_and_act(&gate, &up, self.act)?.to_dtype(out_dtype) } diff --git a/mistralrs-core/src/models/deepseek4.rs b/mistralrs-core/src/models/deepseek4.rs index e2f6f6336..3e84142bf 100644 --- a/mistralrs-core/src/models/deepseek4.rs +++ b/mistralrs-core/src/models/deepseek4.rs @@ -2247,6 +2247,60 @@ impl MoeGate { }) } + /// One-kernel spelling of the `gate.renormalize` chain. + /// `Ok(None)` means the caller must run `eager`. + fn fused_renorm( + &self, + topk_weight: &Tensor, + do_renorm: bool, + eager: &dyn Fn(&Tensor) -> Result, + ) -> Result> { + use crate::cuda::hc_fused as seam; + + if !seam::seam_on("ARC_SEAM_GATE") { + return Ok(None); + } + let decline = |reason: &str| -> Result> { + seam::seam_declined(seam::SEAM_GATE, reason); + Ok(None) + }; + if !topk_weight.device().is_cuda() { + return decline("device is not cuda"); + } + if topk_weight.dtype() != DType::F32 { + return decline("topk_weight is not F32"); + } + let Ok((n, k)) = topk_weight.dims2() else { + return decline("topk_weight is not rank 2"); + }; + if !matches!(k, 1 | 2 | 4 | 6 | 8) { + return decline("top_k outside the specialised set"); + } + if !topk_weight.is_contiguous() { + return decline("topk_weight is not contiguous"); + } + if n == 0 { + return decline("empty gate renorm"); + } + // The scale is narrowed here exactly as candle's `affine` narrows it: + // `T::from_f64(routed_scaling_factor)`. + let out = seam::gate_renorm_cuda( + topk_weight, + 1e-20f64 as f32, + self.cfg.routed_scaling_factor as f32, + do_renorm, + )?; + + if seam::seam_ab_enabled() { + // Already F32 on both sides — nothing is narrowed here, so this + // comparison is the whole contract. + seam::ab_check("seam.gate_renorm", &out, &eager(topk_weight)?)?; + } + + seam::seam_engaged(seam::SEAM_GATE); + Ok(Some(out)) + } + fn forward(&self, xs: &Tensor, input_ids: Option<&Tensor>) -> Result<(Tensor, Tensor)> { let (bs, seq_len, h) = xs.dims3()?; let xs = xs.reshape(((), h))?; @@ -2401,17 +2455,31 @@ impl MoeGate { // `"norm_topk_prob": true` with `sqrtsoftplus`, so this is taken — // identical to the previous scoring-function-derived condition. // Audit finding 21. - if self.cfg.norm_topk_prob + let do_renorm = self.cfg.norm_topk_prob && matches!( self.cfg.scoring_func, ScoringFunc::Sigmoid | ScoringFunc::SqrtSoftplus - ) - { - let denominator = (topk_weight.sum_keepdim(D::Minus1)? + 1e-20)?; - topk_weight = topk_weight.broadcast_div(&denominator)?; - } - - topk_weight = (topk_weight * self.cfg.routed_scaling_factor)?; + ); + // FOUR launches over a [1, 6] tensor — `fast_sum`, `affine(+1e-20)`, + // `bdiv`, `affine(*routed_scaling_factor)` — once per MoE layer per + // token, i.e. 172 per decode forward for 6 divides and 6 multiplies. + // Collapsed to one; `arc_seam_gate_renorm_kernel` in cuda/hc_fused.cu + // replays candle's identity-padded reduction tree and the fmaf that + // `affine` compiles to. `ARC_F32SEAM=0` / `ARC_SEAM_GATE=0` restore the + // chain below. + let eager_renorm = |w: &Tensor| -> Result { + let w = if do_renorm { + let denominator = (w.sum_keepdim(D::Minus1)? + 1e-20)?; + w.broadcast_div(&denominator)? + } else { + w.clone() + }; + w * self.cfg.routed_scaling_factor + }; + topk_weight = match self.fused_renorm(&topk_weight, do_renorm, &eager_renorm)? { + Some(w) => w, + None => eager_renorm(&topk_weight)?, + }; drop(_prof_renorm); // RUN-161 diagnostic: are the selected top-k routing weights peaked or @@ -2685,11 +2753,13 @@ impl Moe { v4_stat_dbg(&y, "moe.routed"); let li = self.layer_idx; v4_collapse_dbg(&y, &format!("L{li}.moe_routed"), 1); - v4_collapse_dbg( - &topk_idx.to_dtype(DType::F32).unwrap_or_else(|_| topk_idx.clone()), - &format!("L{li}.moe_topk_idx"), - 0, - ); + // `to_dtype` in ARGUMENT position is evaluated eagerly, so this debug + // telemetry was casting the U32 expert ids to F32 on every MoE layer of + // every forward — 43 `cast_u32_f32` launches per decoded token — even + // with `ARC_COLLAPSE` unset. `v4_collapse_dbg` already casts to F32 + // internally, *behind* the env guard, so passing the ids through + // unchanged is the same telemetry for zero launches. + v4_collapse_dbg(&topk_idx, &format!("L{li}.moe_topk_idx"), 0); if let Some(ref shared_experts) = self.shared_experts { let _s = arc_profiler::device_span("moe.shared_expert"); @@ -2948,8 +3018,19 @@ fn compressed_row_positions(t_c: usize, ratio: usize, dev: &Device) -> Result cast_u32_f32 -> affine(*ratio) -> + // cast_f32_u32` — three kernel launches to compute a strided range + // that `arange_step` produces directly. The thread-local above makes + // that free on a hit, but a miss (first token, or a window/device + // change) still paid it. The values are identical: the old chain's + // `fmaf(j, ratio, 0.0)` is exact for every `j * ratio` a sequence can + // reach (both operands and the product are integers well inside f32's + // exact range), and `to_dtype(U32)` truncated it back. `ratio` is + // 1 / 4 / 128 (`CompressRatio::ratio`), never 0, so `arange_step`'s + // zero-step bail is unreachable. Multiplied in usize so the product + // cannot wrap before the u32 narrowing. + let t = Tensor::arange_step(0u32, (t_c * ratio) as u32, ratio as u32, dev)?; *c = Some((t_c, ratio, t.clone())); Ok(t) }) diff --git a/mistralrs-core/src/moe/experts.rs b/mistralrs-core/src/moe/experts.rs index ccd3deec2..34d82fea4 100644 --- a/mistralrs-core/src/moe/experts.rs +++ b/mistralrs-core/src/moe/experts.rs @@ -57,6 +57,168 @@ pub(crate) fn swiglu_clamp(gate: &Tensor, up: &Tensor, limit: f32) -> Result<(Te Ok((gate, up)) } +/// The whole clamped-SwiGLU stage — [`swiglu_clamp`], the activation, the +/// multiply and the narrow back to the model dtype — as ONE CUDA kernel. +/// `Ok(None)` means the caller must run the eager chain. +/// +/// # Why this exists +/// +/// The eager spelling is a BF16 -> F32 -> BF16 round trip whose entire middle +/// is elementwise. Per call it costs eight launches (two widening casts, a +/// `minimum`, a `maximum`, a `minimum`, `usilu_f32`, `bmul_f32`, a narrowing +/// cast) — seven for the shared expert, whose activation and multiply are one +/// `fused_glu` — plus **three 4-byte host-to-device memcpys**, because +/// candle's `minimum(f64)`/`maximum(f64)` materialise the bound with +/// `Tensor::new(v, Cpu)?.to_device(cuda)?` on every call +/// (`candle-core/src/tensor.rs`, `binary_op_scalar!`). +/// +/// V4 runs this stage twice per layer (routed experts + shared expert) across +/// 43 layers, so on the measured decode step it is 86 calls: 645 launches and +/// 258 host-to-device copies to compute `silu(clamp(g)) * clamp(u)`. The +/// per-op floor on this H200 is ~1.05 us inside a graph and ~2.6 us on a +/// stream, independent of size, so essentially all of that is launch overhead. +/// `bmaximum_f32 = 86.0/forward` in the kernel histogram is this stage's exact +/// fingerprint — one `maximum` per call, no other producer. +/// +/// # Bit-identity +/// +/// The kernel transcribes the chain rather than re-deriving it; see +/// `cuda/hc_fused.cu`, `arc_seam_swiglu_clamp_kernel`, for the op-by-op +/// citations. The routed and shared spellings differ only in the operand order +/// of the final multiply, which IEEE makes irrelevant. `ARC_SEAM_AB=1` proves +/// it against **both** spellings on live data, comparing at F32 before the +/// narrowing (a BF16 comparison has 8 mantissa bits and would swallow the +/// reassociation errors this is meant to catch); `ARC_HC_AB_POISON=1` is the +/// negative control that must make it fail. +/// +/// Opt out with `ARC_F32SEAM=0` (whole seam) or `ARC_SEAM_SWIGLU=0` (this site). +pub(crate) fn fused_swiglu_clamp( + gate: &Tensor, + up: &Tensor, + limit: f32, + act: Activation, + out_dtype: DType, +) -> Result> { + use crate::cuda::hc_fused as seam; + + if !seam::seam_on("ARC_SEAM_SWIGLU") { + // The A/B leg, not a fallback: stay silent so the DECLINED lines below + // keep meaning "the fast path was wanted and refused". + return Ok(None); + } + // Only Silu is transcribed. `Swish` is `x * sigmoid(x)`, which rounds + // differently from `x / (1 + exp(-x))`, and the rest are not silu at all — + // an unlisted activation must take the eager path, not be approximated. + let refusal = if matches!(act, Activation::Silu) { + swiglu_fusion_refusal(gate, up, matches!(out_dtype, DType::F32 | DType::BF16)) + } else { + Some("activation is not Silu") + }; + if let Some(reason) = refusal { + seam::seam_declined(seam::SEAM_SWIGLU, reason); + return Ok(None); + } + + let out = seam::swiglu_clamp_fused_cuda(gate, up, limit, out_dtype)?; + + if seam::seam_ab_enabled() { + let (g, u) = swiglu_clamp(gate, up, limit)?; + // The routed spelling: `up.mul(&gate.apply(act))`. + let eager_routed = u.mul(&g.apply(&act)?)?; + // The shared spelling: `mul_and_act` -> `fused_glu`, activation cast + // back to T then multiplied in the other order. + let eager_shared = crate::ops::mul_and_act(&g, &u, act)?; + let fused_f32 = seam::swiglu_clamp_fused_cuda(gate, up, limit, DType::F32)?; + seam::ab_check("seam.swiglu.f32.routed", &fused_f32, &eager_routed)?; + seam::ab_check("seam.swiglu.f32.shared", &fused_f32, &eager_shared)?; + seam::ab_check( + "seam.swiglu.out", + &out, + &eager_routed.to_dtype(out_dtype)?, + )?; + } + + seam::seam_engaged(seam::SEAM_SWIGLU); + Ok(Some(out)) +} + +/// Shared common gating for both swiglu fusions. `Err`-free: returns the +/// decline reason, or `None` if the site may fuse. +fn swiglu_fusion_refusal(gate: &Tensor, up: &Tensor, dt_ok: bool) -> Option<&'static str> { + if !gate.device().is_cuda() { + return Some("device is not cuda"); + } + if gate.dtype() != up.dtype() { + return Some("gate and up dtypes differ"); + } + if !matches!(gate.dtype(), DType::F32 | DType::BF16) || !dt_ok { + return Some("dtype outside {F32, BF16}"); + } + if gate.dims() != up.dims() { + return Some("gate and up shapes differ"); + } + if !gate.is_contiguous() || !up.is_contiguous() { + return Some("gate or up is not contiguous"); + } + if gate.elem_count() == 0 { + return Some("empty tensor"); + } + None +} + +/// The **clamp half** of the SwiGLU stage as one kernel, for the shared expert. +/// `Ok(None)` means the caller must run [`swiglu_clamp`]. +/// +/// The shared expert's activation and multiply are +/// `mistralrs_quant::fused_glu`, whose translation unit is compiled with +/// `--use_fast_math` (mistralrs-quant/build.rs:281). Its `expf` is `__expf` +/// and its division is `div.approx`, so the full fusion — which reproduces +/// candle's IEEE `usilu_f32` — would NOT be bit-identical here. Fusing it +/// anyway would be trading the contract for four launches, so only the +/// provable half is taken: two widening casts, `maximum`, two `minimum`s, and +/// the three 4-byte host-to-device copies of the bounds, 5 launches + 3 copies +/// down to 1 launch. `fused_glu` then runs on exactly the F32 tensors it runs +/// on today. +/// +/// Under `ARC_SEAM_AB=1` this ALSO evaluates the full fusion here and reports +/// whether the fast-math and IEEE spellings agree bit for bit — the +/// measurement that would license taking the other four launches. +pub(crate) fn fused_swiglu_clamp_split( + gate: &Tensor, + up: &Tensor, + limit: f32, + act: Activation, +) -> Result> { + use crate::cuda::hc_fused as seam; + + if !seam::seam_on("ARC_SEAM_SWIGLU") { + return Ok(None); + } + if let Some(reason) = swiglu_fusion_refusal(gate, up, true) { + seam::seam_declined(seam::SEAM_SWIGLU_SPLIT, reason); + return Ok(None); + } + + let (g, u) = seam::swiglu_clamp_split_cuda(gate, up, limit)?; + + if seam::seam_ab_enabled() { + let (ge, ue) = swiglu_clamp(gate, up, limit)?; + seam::ab_check("seam.swiglu_split.gate", &g, &ge)?; + seam::ab_check("seam.swiglu_split.up", &u, &ue)?; + // The measurement that decides whether the full fusion is ever legal + // at this site: IEEE silu (this file's kernel) vs the fast-math + // `fused_glu` the shared expert actually runs today. + if matches!(act, Activation::Silu) { + let fused_full = seam::swiglu_clamp_fused_cuda(gate, up, limit, DType::F32)?; + let fastmath = crate::ops::mul_and_act(&ge, &ue, act)?; + seam::ab_check("seam.swiglu_split.ieee_vs_fastmath_glu", &fused_full, &fastmath)?; + } + } + + seam::seam_engaged(seam::SEAM_SWIGLU_SPLIT); + Ok(Some((g, u))) +} + /// Configuration for MoEExperts pub struct MoEExpertsConfig { pub num_experts: usize, @@ -511,6 +673,13 @@ impl MoEExperts { return up.mul(&gate.apply(&self.act)?); }; let out_dtype = gate.dtype(); + // One kernel instead of eight launches and three host-to-device copies. + // The `exp.act_gate` statistic below is only reachable on the eager + // path, which `ARC_F32SEAM=0` restores — the fused kernel never + // materialises the intermediate the statistic reads. + if let Some(out) = fused_swiglu_clamp(gate, up, limit, self.act, out_dtype)? { + return Ok(out); + } let (gate, up) = swiglu_clamp(gate, up, limit)?; let act_gate = gate.apply(&self.act)?; crate::models::deepseek4::v4_stat_dbg(&act_gate, "exp.act_gate"); @@ -794,10 +963,84 @@ impl MoEExperts { }; let _s = arc_profiler::device_span("experts.weighted_sum"); - ys.to_dtype(DType::F32)? - .broadcast_mul(&topk_weights.unsqueeze(D::Minus1)?)? - .sum(D::Minus2)? - .to_dtype(original_dtype) + let eager = || -> Result { + ys.to_dtype(DType::F32)? + .broadcast_mul(&topk_weights.unsqueeze(D::Minus1)?)? + .sum(D::Minus2)? + .to_dtype(original_dtype) + }; + // FOUR launches — widen, weight, reduce over the 6-element expert axis, + // narrow — to compute a 6-term weighted average, 43 times per decoded + // token. Collapsed to one; see `arc_seam_moe_weighted_sum_kernel` in + // cuda/hc_fused.cu, which replays candle's identity-padded `fast_sum` + // tree (K=6 reduces in a block of EIGHT, and the two padding lanes are + // part of the answer). `ARC_F32SEAM=0` / `ARC_SEAM_WSUM=0` restore this. + Ok(match self.fused_weighted_sum(&ys, topk_weights, original_dtype)? { + Some(out) => out, + None => eager()?, + }) + } + + /// One-kernel spelling of the `experts.weighted_sum` chain above. + /// `Ok(None)` means the caller must run the eager chain. + fn fused_weighted_sum( + &self, + ys: &Tensor, + topk_weights: &Tensor, + out_dtype: DType, + ) -> Result> { + use crate::cuda::hc_fused as seam; + + if !seam::seam_on("ARC_SEAM_WSUM") { + return Ok(None); + } + let decline = |reason: &str| -> Result> { + seam::seam_declined(seam::SEAM_WSUM, reason); + Ok(None) + }; + if !ys.device().is_cuda() { + return decline("device is not cuda"); + } + let Ok((n, k, h)) = ys.dims3() else { + return decline("ys is not rank 3"); + }; + if topk_weights.dims() != [n, k] { + return decline("topk_weights shape does not match ys"); + } + // The kernel is specialised on K because the accumulator must live in + // registers; an unlisted K refuses rather than spilling to local memory. + if !matches!(k, 1 | 2 | 4 | 6 | 8) { + return decline("num_experts_per_tok outside the specialised set"); + } + if !matches!(ys.dtype(), DType::F32 | DType::BF16) + || !matches!(out_dtype, DType::F32 | DType::BF16) + { + return decline("dtype outside {F32, BF16}"); + } + if topk_weights.dtype() != DType::F32 { + return decline("topk_weights is not F32"); + } + if !ys.is_contiguous() || !topk_weights.is_contiguous() { + return decline("ys or topk_weights is not contiguous"); + } + if n == 0 || h == 0 { + return decline("empty weighted sum"); + } + + let out = seam::moe_weighted_sum_cuda(ys, topk_weights, out_dtype)?; + + if seam::seam_ab_enabled() { + let eager_f32 = ys + .to_dtype(DType::F32)? + .broadcast_mul(&topk_weights.unsqueeze(D::Minus1)?)? + .sum(D::Minus2)?; + let fused_f32 = seam::moe_weighted_sum_cuda(ys, topk_weights, DType::F32)?; + seam::ab_check("seam.wsum.f32", &fused_f32, &eager_f32)?; + seam::ab_check("seam.wsum.out", &out, &eager_f32.to_dtype(out_dtype)?)?; + } + + seam::seam_engaged(seam::SEAM_WSUM); + Ok(Some(out)) } /// Loop-based forward pass (quantized fallback) diff --git a/mistralrs-core/src/moe/mod.rs b/mistralrs-core/src/moe/mod.rs index 0dad7b9cd..10a990c26 100644 --- a/mistralrs-core/src/moe/mod.rs +++ b/mistralrs-core/src/moe/mod.rs @@ -7,7 +7,7 @@ pub use expert_parallel::{ balancedness, expected_distinct_experts, weight_read_imbalance_bound, Balancedness, BalancednessCounter, ExpertParallelPlan, ExpertPlacement, }; -pub(crate) use experts::swiglu_clamp; +pub(crate) use experts::{fused_swiglu_clamp_split, swiglu_clamp}; pub use experts::{MoEExperts, MoEExpertsConfig}; pub fn shard(dim: usize, rank: usize, world_size: usize) -> Shard { From 671a0c92d441facc7ba052148f62c420a8d13af9 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Thu, 20 Aug 2026 07:41:03 +0100 Subject: [PATCH 26/27] test(ArcGate): per-site A/B accounting, and name the two comparisons that are measurements MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The aggregate A/B counters cannot answer the only question that matters — *which* comparison moved. The first live run made that concrete: 4,222 of 19,000 comparisons mismatched, exactly 2/9, and 2 of the 9 comparisons per layer are deliberately not contracts. Reading that verdict required inferring a ratio, and an inferred verdict is not a verdict. `ab_check` now keeps (comparisons, bad_tensors, bad_elems) per site and dumps the table with the periodic summary, at `warn!` rather than `info!` so the evidence does not depend on the operator having raised RUST_LOG. The two IEEE-silu-vs-fast-math-`fused_glu` comparisons are renamed to `seam.MEASURE.*`. They exist to answer whether the shared-expert site could ever take the full fusion; a difference there is a fact about mistralrs-quant's build flags, not a regression in this change. Co-Authored-By: Claude Opus 5 (1M context) --- mistralrs-core/src/cuda/hc_fused.rs | 32 ++++++++++++++++++++++++++++- mistralrs-core/src/moe/experts.rs | 24 ++++++++++++++++------ 2 files changed, 49 insertions(+), 7 deletions(-) diff --git a/mistralrs-core/src/cuda/hc_fused.rs b/mistralrs-core/src/cuda/hc_fused.rs index 928994ed9..076e903ad 100644 --- a/mistralrs-core/src/cuda/hc_fused.rs +++ b/mistralrs-core/src/cuda/hc_fused.rs @@ -784,6 +784,10 @@ use std::sync::atomic::{AtomicU64, Ordering}; static AB_CMPS: AtomicU64 = AtomicU64::new(0); static AB_BAD_TENSORS: AtomicU64 = AtomicU64::new(0); static AB_BAD_ELEMS: AtomicU64 = AtomicU64::new(0); +/// site -> (comparisons, tensors with any mismatching bit, mismatching elements) +#[allow(clippy::type_complexity)] +static AB_SITES: std::sync::Mutex>> = + std::sync::Mutex::new(None); pub fn ab_enabled() -> bool { static ON: std::sync::OnceLock = std::sync::OnceLock::new(); @@ -851,6 +855,22 @@ pub fn ab_check( } let n = AB_CMPS.fetch_add(1, Ordering::Relaxed) + 1; + // PER-SITE accounting. The aggregate counters alone cannot answer the only + // question that matters — *which* comparison moved — and this file + // deliberately runs two comparisons that are MEASUREMENTS rather than + // contracts (IEEE silu vs the fast-math `fused_glu` the shared expert + // actually runs). Without a per-site table the verdict has to be inferred + // from a ratio, and an inferred verdict is not a verdict. + { + let mut g = AB_SITES.lock().unwrap_or_else(|e| e.into_inner()); + let m = g.get_or_insert_with(std::collections::BTreeMap::new); + let e = m.entry(site.to_string()).or_insert((0u64, 0u64, 0u64)); + e.0 += 1; + if bad > 0 { + e.1 += 1; + e.2 += bad; + } + } if bad > 0 { AB_BAD_TENSORS.fetch_add(1, Ordering::Relaxed); AB_BAD_ELEMS.fetch_add(bad, Ordering::Relaxed); @@ -866,13 +886,23 @@ pub fn ab_check( } } if n % 1000 == 0 { - tracing::info!( + // `warn!`, not `info!`: the summary is the evidence, and it must not + // depend on the operator having remembered to raise RUST_LOG. + tracing::warn!( "HC_AB: {n} comparisons, {} tensors with mismatches, {} bad elements \ (poison={})", AB_BAD_TENSORS.load(Ordering::Relaxed), AB_BAD_ELEMS.load(Ordering::Relaxed), ab_poison() ); + let g = AB_SITES.lock().unwrap_or_else(|e| e.into_inner()); + if let Some(m) = g.as_ref() { + for (site, (c, bt, be)) in m.iter() { + tracing::warn!( + "HC_AB_SITE {site} comparisons={c} bad_tensors={bt} bad_elems={be}" + ); + } + } } Ok(bad) } diff --git a/mistralrs-core/src/moe/experts.rs b/mistralrs-core/src/moe/experts.rs index 34d82fea4..b75a6691a 100644 --- a/mistralrs-core/src/moe/experts.rs +++ b/mistralrs-core/src/moe/experts.rs @@ -129,12 +129,20 @@ pub(crate) fn fused_swiglu_clamp( // back to T then multiplied in the other order. let eager_shared = crate::ops::mul_and_act(&g, &u, act)?; let fused_f32 = seam::swiglu_clamp_fused_cuda(gate, up, limit, DType::F32)?; - seam::ab_check("seam.swiglu.f32.routed", &fused_f32, &eager_routed)?; - seam::ab_check("seam.swiglu.f32.shared", &fused_f32, &eager_shared)?; + // CONTRACT: this site replaces the routed chain, so it must match it. + seam::ab_check("seam.swiglu.f32", &fused_f32, &eager_routed)?; + seam::ab_check("seam.swiglu.out", &out, &eager_routed.to_dtype(out_dtype)?)?; + // MEASUREMENT, not contract: the same value spelled with the fast-math + // `fused_glu` (mistralrs-quant is built with `--use_fast_math`). The + // routed site never used that spelling, so a difference here is not a + // regression — it is the number that says whether the shared-expert + // site could ever take the full fusion. Named `MEASURE.` so the verdict + // script can tell a measurement from a contract instead of failing the + // whole leg on it. seam::ab_check( - "seam.swiglu.out", - &out, - &eager_routed.to_dtype(out_dtype)?, + "seam.MEASURE.ieee_silu_vs_fastmath_glu.at_routed", + &fused_f32, + &eager_shared, )?; } @@ -211,7 +219,11 @@ pub(crate) fn fused_swiglu_clamp_split( if matches!(act, Activation::Silu) { let fused_full = seam::swiglu_clamp_fused_cuda(gate, up, limit, DType::F32)?; let fastmath = crate::ops::mul_and_act(&ge, &ue, act)?; - seam::ab_check("seam.swiglu_split.ieee_vs_fastmath_glu", &fused_full, &fastmath)?; + seam::ab_check( + "seam.MEASURE.ieee_silu_vs_fastmath_glu.at_shared", + &fused_full, + &fastmath, + )?; } } From b39a206db82067ceec7c54c15aaa408583b5f1e1 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Thu, 20 Aug 2026 08:03:31 +0100 Subject: [PATCH 27/27] docs(ArcInfer/ArcMoE): record the measured F32-seam result MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Measured on arc-graph1 (H200), ONE binary toggled with ARC_F32SEAM, 160-token greedy generation, every count a full GROUP BY over the nsys sqlite, step count pinned by two independent anchors that agree at 160 forwards. per forward ARC_F32SEAM=0 =1 delta kernel launches 4,622.5 3,891.5 -731.0 (-15.8%) cuMemcpyHtoDAsync 1,450.9 762.9 -688.0 (-47.4%) ALL host-issued ops 8,081.8 6,645.4 -1,436.4 (-17.8%) Decode, interleaved A-B-A-B: 28.79 / 31.84 / 28.04 / 31.65 tok/s. Controls agree to 2.7%, treatments to 0.6%, no overlap. 35.20 -> 31.51 ms/token, 28.42 -> 31.75 tok/s (+11.7%). SM clock 1830 MHz during the measured request in all four legs. 320 greedy tokens byte-identical across the toggle (sha 6fec510d6ef723cb, 1,141 chars, all four legs). Bit-identity: 14,778 comparisons over seven contract sites, zero mismatching bits, with the 1-ULP negative control firing at every one of them. Two findings worth their own lines: The 688 vanished host-to-device copies are not only the 258 clamp bounds. candle uploads a dims/strides `info` array with every strided op, so each removed launch took its bookkeeping copy with it. HtoD copies are 1,451 per forward on this model and average 57 bytes; they are op count wearing a different hat. Both `seam.MEASURE.ieee_silu_vs_fastmath_glu` sites differ on 2,111/2,111 tensors, ~38% of elements each. That is the number that says the shared expert can NOT take the full fusion while mistralrs-quant builds `fused_glu` with `--use_fast_math` — the decision to fuse only its clamp half is now measured rather than argued. Co-Authored-By: Claude Opus 5 (1M context) --- mistralrs-core/src/cuda/hc_fused.cu | 41 +++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/mistralrs-core/src/cuda/hc_fused.cu b/mistralrs-core/src/cuda/hc_fused.cu index 29f0c629c..19fe63b96 100644 --- a/mistralrs-core/src/cuda/hc_fused.cu +++ b/mistralrs-core/src/cuda/hc_fused.cu @@ -529,6 +529,47 @@ __global__ void hc_post_fused_kernel( // matter. `ARC_HC_AB_POISON=1` perturbs one element by 1 ULP and must make // every comparison fail; if it does not, the comparison is broken, not the // kernel. +// +// --------------------------------------------------------------------------- +// MEASURED RESULT -- arc-graph1 (H200), agent/f32-cast-fuse @ 9de64af, +// ONE binary toggled with ARC_F32SEAM, 160-token greedy generation, nsys +// `-t cuda`, every count a full GROUP BY over the exported sqlite, step count +// pinned by two independent anchors that agree at 160 forwards. +// +// per forward ARC_F32SEAM=0 =1 delta +// kernel launches 4,622.5 3,891.5 -731.0 (-15.8%) +// cuMemcpyHtoDAsync 1,450.9 762.9 -688.0 (-47.4%) +// cuMemAllocAsync 205.4 196.7 -8.7 +// ALL host-issued ops 8,081.8 6,645.4 -1,436.4 (-17.8%) +// +// Per kernel, and every line is the arithmetic this file predicted: +// bminimum_f32 172.0 -> 0.0 bmaximum_f32 86.0 -> 0.0 +// usilu_f32 43.0 -> 0.0 bdiv_f32 48.4 -> 5.4 +// bmul_f32 94.4 -> 8.4 fast_sum_f32 98.8 -> 12.8 +// cast_bf16_f32 356.8 -> 141.8 cast_f32_bf16 221.1 -> 135.1 +// affine_f32 185.0 -> 99.0 + four new kernels at 43.0 each +// +// The 688 vanished host-to-device copies are NOT only the clamp bounds (258 of +// them). candle uploads a dims/strides `info` array with every strided op, so +// each removed launch took its bookkeeping copy with it -- 83.1 kB/forward of +// ~57-byte uploads down to 54.9 kB. +// +// Decode, interleaved A-B-A-B because a single before/after on a shared box is +// not evidence: 28.79 / 31.84 / 28.04 / 31.65 tok/s. The two controls agree to +// 2.7%, the two treatments to 0.6%, and neither pair overlaps the other: +// 35.20 -> 31.51 ms/token, 28.42 -> 31.75 tok/s (+11.7%). SM clock read +// 1830 MHz DURING the measured request in all four legs (1830 is this card's +// real load clock; sampling around a short warm reads the idle clock instead). +// +// Output identity: 320 greedy tokens, 1,141 characters, sha 6fec510d6ef723cb +// in ALL FOUR legs -- byte-identical across the toggle. +// +// Bit-identity: 14,778 comparisons across the seven contract sites, ZERO +// mismatching bits. The negative control fires at every one of them +// (2,111/2,111 tensors, exactly one bad element each -- it detects a +// single-element 1-ULP perturbation everywhere), so the clean leg is a result +// rather than a silence. +// --------------------------------------------------------------------------- // V4 clamped SwiGLU: 8 launches + 3 H2D memcpys -> 1 kernel. //