From a6ca9fba4fa6e89a82c15d69bcbe8edff83d289f Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 23 Apr 2026 10:07:06 +0200 Subject: [PATCH] feat(mamba): wire real GPU Bernoulli dropout in supervised Mamba2 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Supervised Mamba2's dropout_rate field was configured and stored but never applied — comments said "simulated" but the arithmetic was absent. An operator tuning dropout_rate upward got zero regularisation. New dropout_kernel.cu with a forward-only Bernoulli dropout (Philox- seeded, deterministic, in-place). New build.rs matching ml-dqn's pattern. Applied in forward_with_gradients (training path) only; the `forward()` path used by validate/predict/SPSA stays deterministic so SPSA's ±ε finite-difference estimator is not destabilised. NOT a DQN fix: DQN has its own regime_dropout kernel in cuda_pipeline/experience_kernels.cu and its own native mamba2_step in gpu_dqn_trainer.rs; DQN does not call into Mamba2SSM. This commit affects only the supervised Mamba2 trainer and its hyperopt adapter. Tests: determinism, eval-mode identity, ctr-increment divergence. DQN smoke tests (magnitude_distribution, multi_fold_convergence) are unaffected by this change — ran them for regression assurance only. Co-Authored-By: Claude Opus 4.7 (1M context) --- crates/ml-supervised/build.rs | 149 ++++++++ .../ml-supervised/src/mamba/dropout_kernel.cu | 71 ++++ crates/ml-supervised/src/mamba/mod.rs | 341 +++++++++++++++++- 3 files changed, 555 insertions(+), 6 deletions(-) create mode 100644 crates/ml-supervised/build.rs create mode 100644 crates/ml-supervised/src/mamba/dropout_kernel.cu diff --git a/crates/ml-supervised/build.rs b/crates/ml-supervised/build.rs new file mode 100644 index 000000000..dd1f05508 --- /dev/null +++ b/crates/ml-supervised/build.rs @@ -0,0 +1,149 @@ +//! Precompile CUDA kernels owned by `ml-supervised`. +//! +//! Modelled on `crates/ml-dqn/build.rs` (the canonical reference). Emits +//! `.cubin` files into `$OUT_DIR`; Rust callers `include_bytes!` them via +//! `concat!(env!("OUT_DIR"), "/.cubin")` and load through +//! `stream.context().load_cubin(...)`. +//! +//! Kernels compiled here: +//! - `mamba/dropout_kernel.cu` → `dropout_kernel.cubin` (forward-only +//! Bernoulli inverted dropout used by `Mamba2SSM::forward_with_gradients`). + +use std::path::{Path, PathBuf}; +use std::process::Command; + +fn main() { + println!("cargo:rerun-if-changed=build.rs"); + + // Only compile CUDA kernels when the cuda feature is enabled (matches + // ml-dqn's gating exactly). Non-CUDA builds skip kernel compilation. + if std::env::var("CARGO_FEATURE_CUDA").is_err() { + return; + } + + let out_dir = PathBuf::from(std::env::var("OUT_DIR").unwrap()); + let manifest_dir = PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").unwrap()); + + // Detect GPU architecture from env or default to sm_80 (L40S / A100). + let cuda_compute_cap = std::env::var("CUDA_COMPUTE_CAP").unwrap_or_else(|_| "80".to_string()); + let arch = format!("sm_{cuda_compute_cap}"); + + // Check if nvcc is available. + let nvcc = match find_nvcc() { + Some(p) => p, + None => { + eprintln!(" warning: nvcc not found, skipping ml-supervised CUDA kernel precompilation"); + eprintln!(" Install CUDA toolkit or set CUDA_HOME for GPU builds"); + return; + } + }; + + // Kernels owned by ml-supervised. Each entry is (source path relative to + // crate root, output cubin base name). The kernel is self-contained and + // does NOT consume ml/src/cuda_pipeline/common_device_functions.cuh — + // that header carries DQN-specific constants (DQN_NUM_ACTIONS, + // STATE_DIM, ...) that are irrelevant here. + let kernels: &[(&str, &str)] = &[ + ("src/mamba/dropout_kernel.cu", "dropout_kernel.cu"), + ]; + + let mut failed: Vec<&str> = Vec::new(); + + for (src_path, kernel_name) in kernels { + let full_src = manifest_dir.join(src_path); + if !try_compile_kernel(&nvcc, &full_src, kernel_name, &arch, &out_dir) { + failed.push(kernel_name); + } + } + + let total = kernels.len(); + let passed = total - failed.len(); + eprintln!( + " ml-supervised: Precompiled {passed}/{total} CUDA kernels ({arch})", + ); + if !failed.is_empty() { + eprintln!(" FAILED: {}", failed.join(", ")); + panic!( + "nvcc failed to compile {} ml-supervised kernel(s): {}", + failed.len(), + failed.join(", ") + ); + } +} + +/// Compile a single standalone `.cu` kernel to a `.cubin` via nvcc. +/// +/// No header prepending — the kernel is expected to be self-contained. +fn try_compile_kernel( + nvcc: &Path, + kernel_path: &Path, + kernel_name: &str, + arch: &str, + out_dir: &Path, +) -> bool { + println!("cargo:rerun-if-changed={}", kernel_path.display()); + + if !kernel_path.exists() { + eprintln!(" FAILED: {kernel_name} (source not found at {})", kernel_path.display()); + return false; + } + + let cubin_name = kernel_name.replace(".cu", ".cubin"); + let cubin_path = out_dir.join(&cubin_name); + + let status = Command::new(nvcc) + .args([ + "-cubin", + &format!("-arch={arch}"), + "-O3", + "--ftz=true", + "--fmad=true", + "--prec-div=true", + "--prec-sqrt=true", + "-o", + cubin_path.to_str().unwrap(), + kernel_path.to_str().unwrap(), + ]) + .status(); + + match status { + Ok(s) if s.success() => { + eprintln!(" Compiled {kernel_name} -> {cubin_name} ({arch})"); + true + } + Ok(s) => { + eprintln!( + " FAILED: {kernel_name} (exit={})", + s.code().unwrap_or(-1) + ); + false + } + Err(e) => { + eprintln!(" FAILED: {kernel_name} (nvcc error: {e})"); + false + } + } +} + +/// Find nvcc: prefer `$CUDA_HOME/bin/nvcc`, then check common install paths, +/// then PATH. +fn find_nvcc() -> Option { + if let Ok(home) = std::env::var("CUDA_HOME") { + let nvcc = PathBuf::from(home).join("bin/nvcc"); + if nvcc.exists() { + return Some(nvcc); + } + } + + for path in &["/usr/local/cuda/bin/nvcc", "/usr/bin/nvcc"] { + let p = PathBuf::from(path); + if p.exists() { + return Some(p); + } + } + + match Command::new("nvcc").arg("--version").output() { + Ok(output) if output.status.success() => Some(PathBuf::from("nvcc")), + _ => None, + } +} diff --git a/crates/ml-supervised/src/mamba/dropout_kernel.cu b/crates/ml-supervised/src/mamba/dropout_kernel.cu new file mode 100644 index 000000000..ef23021c8 --- /dev/null +++ b/crates/ml-supervised/src/mamba/dropout_kernel.cu @@ -0,0 +1,71 @@ +// Bernoulli inverted dropout — forward-only, deterministic, in-place. +// +// Philox-style stateless hash parameterised by (seed, ctr, layer_idx, i) so +// the same (seed, ctr, layer_idx) always yields the same mask, and distinct +// layers never alias (they carry different layer_idx into the hash tuple). +// +// Inverted dropout: surviving elements are scaled by 1/keep_prob so the +// expected value matches the eval-time (dropout-off) activation. +// +// Used only by `Mamba2SSM::forward_with_gradients` (training path). The pure +// `forward()` used by validate / predict_single_fast / SPSA's ±ε probes is +// left untouched so SPSA's finite-difference estimator stays stable. + +/* Stateless Philox-like mixing hash. Same 32-bit avalanche used by the DQN + * experience kernels' philox_uniform, extended to fold in a 64-bit seed, + * 64-bit counter, layer index, and element index. Returns a uniform f32 in + * [0, 1). */ +__device__ __forceinline__ float mamba_dropout_uniform( + unsigned int seed_lo, + unsigned int seed_hi, + unsigned int ctr_lo, + unsigned int ctr_hi, + unsigned int layer_idx, + unsigned int i +) { + unsigned int key = seed_lo ^ (seed_hi * 0x9E3779B9u) ^ (layer_idx * 0x85EBCA6Bu); + unsigned int ctr = ctr_lo + i; + ctr ^= ctr_hi * 0xC2B2AE35u; + ctr ^= key * 0x9E3779B9u; + ctr *= 0x85EBCA6Bu; + ctr ^= ctr >> 13; + ctr *= 0xC2B2AE35u; + ctr ^= ctr >> 16; + return (float)(ctr & 0x00FFFFFFu) / 16777216.0f; +} + +/* Forward-only Bernoulli inverted dropout applied in-place. + * + * out[i] = in[i] * (u_i < keep_prob ? 1/keep_prob : 0) + * + * where u_i = mamba_dropout_uniform(seed, ctr, layer_idx, i). + * + * Grid: ceil(n / block_dim.x), Block: typically 256. + * + * Args: + * data - device pointer to n f32 elements (read/write, in-place). + * n - element count. + * keep_prob - 1 - dropout_rate, in (0, 1]. + * seed_lo - low 32 bits of the 64-bit dropout seed. + * seed_hi - high 32 bits of the 64-bit dropout seed. + * ctr_lo - low 32 bits of the 64-bit dropout counter. + * ctr_hi - high 32 bits of the 64-bit dropout counter. + * layer_idx - layer index (0..num_layers). + */ +extern "C" __global__ void bernoulli_dropout_apply( + float* __restrict__ data, + int n, + float keep_prob, + unsigned int seed_lo, + unsigned int seed_hi, + unsigned int ctr_lo, + unsigned int ctr_hi, + unsigned int layer_idx +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= n) return; + + float u = mamba_dropout_uniform(seed_lo, seed_hi, ctr_lo, ctr_hi, layer_idx, (unsigned int)idx); + float mask = (u < keep_prob) ? (1.0f / keep_prob) : 0.0f; + data[idx] = data[idx] * mask; +} diff --git a/crates/ml-supervised/src/mamba/mod.rs b/crates/ml-supervised/src/mamba/mod.rs index fb1b40b4a..4820cb224 100644 --- a/crates/ml-supervised/src/mamba/mod.rs +++ b/crates/ml-supervised/src/mamba/mod.rs @@ -57,7 +57,7 @@ use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant, SystemTime}; -use cudarc::driver::CudaStream; +use cudarc::driver::{CudaFunction, CudaStream, LaunchConfig, PushKernelArg}; use serde::{Deserialize, Serialize}; use tracing::{debug, info, instrument, trace, warn}; use uuid::Uuid; @@ -91,6 +91,13 @@ const fn default_spsa_epsilon() -> f64 { 0.01 } +/// Default Philox seed for the forward-only Bernoulli dropout kernel in +/// `forward_with_gradients`. Callers needing deterministic replay (e.g. +/// tests) should use [`Mamba2SSM::new_with_dropout_seed`] to pin the seed +/// explicitly; the struct-level counter starts at 0 and advances by 1 per +/// `forward_with_gradients` call. +pub const DEFAULT_DROPOUT_SEED: u64 = 0xDA7A_FACE_DA7A_FACE; + /// Configuration for `MAMBA-2` state-space model #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Mamba2Config { @@ -474,6 +481,30 @@ impl CudaLayerNorm { } } +/// Load the forward-only Bernoulli dropout kernel from the cubin precompiled +/// by `build.rs`. Returns `None` (not an error) if the cubin lookup or +/// function load fails at runtime — call sites must handle the no-op +/// fallback explicitly rather than panicking, matching the ml-dqn pattern of +/// degrading gracefully when GPU support is absent. +/// +/// On non-cuda builds the compiled cubin does not exist (build.rs early +/// returns), so the whole function body is gated on `feature = "cuda"` and +/// the non-cuda branch always yields `None`. +#[cfg(feature = "cuda")] +fn load_dropout_kernel(stream: &Arc) -> Option { + static CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/dropout_kernel.cubin")); + let module = stream + .context() + .load_cubin(CUBIN.to_vec()) + .ok()?; + module.load_function("bernoulli_dropout_apply").ok() +} + +#[cfg(not(feature = "cuda"))] +fn load_dropout_kernel(_stream: &Arc) -> Option { + None +} + /// `MAMBA-2` State-Space Model implementation pub struct Mamba2SSM { pub config: Mamba2Config, @@ -492,6 +523,18 @@ pub struct Mamba2SSM { pub layer_norms: Vec, pub dropout_rate: f32, + /// Compiled forward-only Bernoulli dropout kernel. `None` on non-CUDA + /// builds (feature `cuda` disabled) or if precompilation skipped because + /// nvcc was unavailable at build time — callers fall back to a no-op. + dropout_kernel: Option, + /// 64-bit Philox seed. Default: `0xDA7A_FACE_DA7A_FACE`; override via + /// [`Mamba2SSM::new_with_dropout_seed`] for deterministic replay tests. + dropout_seed: u64, + /// Monotonic counter advanced once per `forward_with_gradients` call so + /// successive training steps see independent dropout masks even with the + /// same seed. Fed to the Philox hash alongside `(layer_idx, element_idx)`. + dropout_ctr: u64, + // Training state pub optimizer_state: HashMap, pub gradients: HashMap, @@ -524,7 +567,11 @@ impl Mamba2SSM { GpuTensor::from_vec(vec![value as f32], &[1], stream) } - /// Create new `MAMBA-2` model + /// Create new `MAMBA-2` model with the default dropout seed + /// (`0xDA7A_FACE_DA7A_FACE`). + /// + /// Delegates to [`Mamba2SSM::new_with_dropout_seed`]. Use that variant + /// directly when a test needs a specific deterministic mask stream. /// /// # Errors /// @@ -534,6 +581,28 @@ impl Mamba2SSM { /// - Layer norm creation fails /// - SSD layer initialization fails pub fn new(config: Mamba2Config, stream: &Arc) -> Result { + Self::new_with_dropout_seed(config, stream, DEFAULT_DROPOUT_SEED) + } + + /// Create new `MAMBA-2` model with an explicit 64-bit dropout seed. + /// + /// The seed feeds the Philox hash used by the forward-only Bernoulli + /// dropout kernel in `forward_with_gradients`. Two models with the same + /// `(dropout_seed, dropout_ctr)` produce bit-identical masks for the same + /// input shapes — used by `test_dropout_determinism`. + /// + /// # Errors + /// + /// Returns `MLError` if: + /// - CUDA stream creation fails + /// - Linear layer creation fails + /// - Layer norm creation fails + /// - SSD layer initialization fails + pub fn new_with_dropout_seed( + config: Mamba2Config, + stream: &Arc, + dropout_seed: u64, + ) -> Result { if config.d_model == 0 { return Err(MLError::ConfigError("Mamba2 requires d_model > 0".to_owned())); } @@ -586,6 +655,13 @@ impl Mamba2SSM { // Store learning_rate before moving config let learning_rate = config.learning_rate; + // Load the forward-only Bernoulli dropout kernel precompiled by + // build.rs. A missing kernel (non-cuda build, or nvcc unavailable at + // build time) degrades to a no-op in `apply_dropout_inplace` rather + // than panicking — the ghost-feature repair only fires where the + // precompiled kernel is available. + let dropout_kernel = load_dropout_kernel(stream); + Ok(Self { config, metadata, @@ -600,6 +676,9 @@ impl Mamba2SSM { output_projection, layer_norms, dropout_rate, + dropout_kernel, + dropout_seed, + dropout_ctr: 0, optimizer_state: HashMap::new(), gradients: HashMap::new(), grad_scaler: 1.0, @@ -693,8 +772,11 @@ impl Mamba2SSM { // Residual connection hidden = gpu_add(&hidden, &layer_output)?; - // Dropout (simulated — multiply by (1 - dropout_rate) during training) - // TODO: proper GPU dropout kernel + // Dropout applied only in `forward_with_gradients` (training + // path); inference and SPSA must stay deterministic. SPSA's + // ±ε probes in `backward_pass` call this function twice, and + // asymmetric masks across those two calls would poison the + // finite-difference estimator — keep this path mask-free. } // Output projection with sigmoid activation (P0 FIX: bound output to [0,1] for normalized targets) @@ -1470,8 +1552,98 @@ impl Mamba2SSM { Ok(loss_value) } - /// Forward pass with gradient computation enabled + /// In-place Bernoulli inverted dropout on the post-residual hidden state. + /// + /// Forward-only. Launches the precompiled `bernoulli_dropout_apply` + /// kernel with Philox inputs `(dropout_seed, dropout_ctr, layer_idx, i)` + /// so the mask is deterministic under `(seed, ctr, layer_idx)` and + /// non-aliasing across layers. The counter is advanced once per + /// `forward_with_gradients` call, not per layer, so a full pass sees + /// distinct masks per layer via `layer_idx` while successive training + /// steps see distinct masks via `dropout_ctr`. + /// + /// No-ops when: + /// - `dropout_rate <= 0.0` (eval-mode identity, byte-for-byte), + /// - tensor is empty, + /// - the precompiled kernel is unavailable (non-cuda build path). + /// + /// Clamps `dropout_rate` to `< 1.0` to avoid a divide-by-zero keep_prob. + #[allow(unsafe_code)] // CUDA kernel launch via cudarc. + pub(crate) fn apply_dropout_inplace( + &self, + hidden: &mut GpuTensor, + layer_idx: usize, + ) -> Result<(), MLError> { + if self.dropout_rate <= 0.0 { + return Ok(()); + } + let n = hidden.numel(); + if n == 0 { + return Ok(()); + } + let Some(kernel) = self.dropout_kernel.as_ref() else { + // Kernel missing (non-cuda build / nvcc-less build) — soft-stub + // by doing nothing. Logged once per call so regressions surface + // in training output rather than silently disappearing. + tracing::debug!( + "apply_dropout_inplace: dropout kernel unavailable, skipping \ + (dropout_rate={} layer_idx={})", + self.dropout_rate, + layer_idx, + ); + return Ok(()); + }; + + // Clamp to avoid divide-by-zero on keep_prob. + let drop_rate = self.dropout_rate.min(1.0 - 1e-6).max(0.0); + let keep_prob = 1.0_f32 - drop_rate; + + let n_i32 = n as i32; + let seed_lo = (self.dropout_seed & 0xFFFF_FFFF) as u32; + let seed_hi = (self.dropout_seed >> 32) as u32; + let ctr_lo = (self.dropout_ctr & 0xFFFF_FFFF) as u32; + let ctr_hi = (self.dropout_ctr >> 32) as u32; + let layer_u32 = layer_idx as u32; + + let block_dim: u32 = 256; + let grid_dim: u32 = ((n as u32) + block_dim - 1) / block_dim; + let cfg = LaunchConfig { + grid_dim: (grid_dim, 1, 1), + block_dim: (block_dim, 1, 1), + shared_mem_bytes: 0, + }; + + // SAFETY: hidden.data is a valid device allocation on self.stream. + // n_i32 matches the allocation length. All scalar args are by-value. + unsafe { + self.stream + .launch_builder(kernel) + .arg(&mut hidden.data) + .arg(&n_i32) + .arg(&keep_prob) + .arg(&seed_lo) + .arg(&seed_hi) + .arg(&ctr_lo) + .arg(&ctr_hi) + .arg(&layer_u32) + .launch(cfg) + .map_err(|e| MLError::ModelError(format!("bernoulli_dropout_apply: {e}")))?; + } + + Ok(()) + } + + /// Forward pass with gradient computation enabled. + /// + /// This is the sole training-path forward. Dropout fires here only; the + /// pure `forward()` used by `validate()`, `predict_single_fast()`, and + /// SPSA's ±ε probes in `backward_pass()` stays deterministic so SPSA's + /// finite-difference estimator is not destabilised by asymmetric masks. pub fn forward_with_gradients(&mut self, input: &GpuTensor) -> Result { + // Advance the Philox counter once per call. wrapping_add avoids UB + // on overflow after ~1.8e19 training steps. + self.dropout_ctr = self.dropout_ctr.wrapping_add(1); + // Input projection with gradients let mut hidden = self.input_projection.forward(input)?; @@ -1494,7 +1666,11 @@ impl Mamba2SSM { // Residual connection hidden = gpu_add(&hidden, &layer_output)?; - // Dropout (simulated — TODO: proper GPU dropout kernel) + // Bernoulli inverted dropout (GPU kernel, Philox-seeded, in-place). + // Previously this site carried a "simulated — TODO" comment but no + // actual op — operators tuning dropout_rate got zero effect. See + // `apply_dropout_inplace` for the kernel launch. + self.apply_dropout_inplace(&mut hidden, layer_idx)?; } // Output projection @@ -3367,3 +3543,156 @@ mod extra_tests { bilinear, exact, zoh); } } + +/// Tests for the forward-only Bernoulli dropout kernel wired into +/// `Mamba2SSM`. Gated on `feature = "cuda"` and `#[ignore]` so they only +/// run when explicitly requested on a CUDA-capable host (matches the +/// `magnitude_distribution` / `multi_fold_convergence` smoke-test pattern). +#[cfg(all(test, feature = "cuda"))] +#[allow(clippy::unwrap_used, clippy::float_cmp, clippy::cast_precision_loss)] +mod dropout_tests { + use super::*; + use cudarc::driver::CudaContext; + + /// Minimal config: dropout machinery only needs a stream + one layer + /// to own a layer_idx. Numeric model parameters are unused here. + fn minimal_config(dropout: f64) -> Mamba2Config { + Mamba2Config { + d_model: 4, + d_state: 2, + d_head: 2, + num_heads: 2, + expand: 1, + num_layers: 1, + dropout, + use_ssd: false, + use_selective_state: false, + hardware_aware: false, + batch_size: 1, + seq_len: 4, + ..Default::default() + } + } + + fn make_model(dropout: f64, seed: u64) -> Mamba2SSM { + let ctx = CudaContext::new(0).expect("CUDA context required"); + let stream = ctx.new_stream().expect("Failed to create CUDA stream"); + Mamba2SSM::new_with_dropout_seed(minimal_config(dropout), &stream, seed) + .expect("construct Mamba2SSM") + } + + fn make_input(model: &Mamba2SSM, n: usize) -> GpuTensor { + // Deterministic non-zero values so a 0-mask is unambiguous. + let values: Vec = (0..n).map(|i| 1.0 + (i as f32) * 0.01).collect(); + GpuTensor::from_vec(values, &[n], &model.stream).expect("upload input") + } + + /// Two calls with the same `(dropout_seed, dropout_ctr)` must produce + /// bit-identical output. Advancing the ctr on each side by the same + /// amount preserves that equivalence. + #[test] + #[ignore = "requires CUDA hardware"] + fn test_dropout_determinism() { + let n = 1024; + let mut model_a = make_model(0.3, 0x1234_5678_9ABC_DEF0); + let mut model_b = make_model(0.3, 0x1234_5678_9ABC_DEF0); + + // Both models share ctr=0. Advance each to ctr=1 uniformly. + model_a.dropout_ctr = 1; + model_b.dropout_ctr = 1; + + let mut tensor_a = make_input(&model_a, n); + let mut tensor_b = make_input(&model_b, n); + + model_a.apply_dropout_inplace(&mut tensor_a, 0).unwrap(); + model_b.apply_dropout_inplace(&mut tensor_b, 0).unwrap(); + + let host_a = tensor_a.to_vec().unwrap(); + let host_b = tensor_b.to_vec().unwrap(); + assert_eq!(host_a.len(), n); + assert_eq!(host_b.len(), n); + for (i, (a, b)) in host_a.iter().zip(host_b.iter()).enumerate() { + assert_eq!( + a.to_bits(), + b.to_bits(), + "dropout mask diverged at i={i}: a={a} b={b} — Philox not deterministic under (seed, ctr, layer_idx)" + ); + } + + // Sanity: at drop_rate=0.3 with n=1024 we expect ~307 zeros. A + // functioning kernel should drop at least a handful; zero zeros + // would indicate the mask arithmetic is absent again. + let zero_count = host_a.iter().filter(|&&v| v == 0.0).count(); + assert!( + zero_count > 0, + "dropout_rate=0.3 produced zero masked elements over {n} samples \ + — kernel is not applying the mask" + ); + } + + /// With `dropout_rate = 0.0` the kernel must short-circuit to a no-op: + /// output bytes must match input bytes exactly. This is the eval-mode + /// identity contract — a non-zero dropout_rate with a zero-valued mask + /// would scale surviving elements by 1/keep_prob and violate identity. + #[test] + #[ignore = "requires CUDA hardware"] + fn test_dropout_eval_mode_identity() { + let n = 512; + let mut model = make_model(0.0, DEFAULT_DROPOUT_SEED); + model.dropout_ctr = 42; + + let input_values: Vec = (0..n).map(|i| 1.0 + (i as f32) * 0.01).collect(); + let mut tensor = GpuTensor::from_vec(input_values.clone(), &[n], &model.stream) + .expect("upload input"); + + model.apply_dropout_inplace(&mut tensor, 0).unwrap(); + + let host = tensor.to_vec().unwrap(); + assert_eq!(host.len(), n); + for (i, (before, after)) in input_values.iter().zip(host.iter()).enumerate() { + assert_eq!( + before.to_bits(), + after.to_bits(), + "dropout_rate=0 modified element i={i}: before={before} after={after} \ + — identity short-circuit broken" + ); + } + } + + /// Two calls with distinct counters (same seed) must diverge. Otherwise + /// the training loop would see the same mask every step, collapsing the + /// regularisation to a single random subset. + #[test] + #[ignore = "requires CUDA hardware"] + fn test_dropout_counter_increment_diverges() { + let n = 1024; + let mut model = make_model(0.3, 0xCAFE_F00D_CAFE_F00D); + + model.dropout_ctr = 1; + let mut tensor_step1 = make_input(&model, n); + model.apply_dropout_inplace(&mut tensor_step1, 0).unwrap(); + let host_step1 = tensor_step1.to_vec().unwrap(); + + model.dropout_ctr = 2; + let mut tensor_step2 = make_input(&model, n); + model.apply_dropout_inplace(&mut tensor_step2, 0).unwrap(); + let host_step2 = tensor_step2.to_vec().unwrap(); + + let diverged = host_step1 + .iter() + .zip(host_step2.iter()) + .filter(|(a, b)| a.to_bits() != b.to_bits()) + .count(); + + // At drop_rate=0.3 over 1024 samples we'd expect ~40% of positions + // to disagree between two independent masks (either zeroed vs kept, + // or mixed scaling). Require at least 10% divergence — a single + // digit count would mean the counter is not being mixed into the + // Philox tuple. + assert!( + diverged > n / 10, + "counter increment produced only {diverged}/{n} diverged elements — \ + dropout_ctr not folded into Philox tuple" + ); + } +}