From c7bb79a6250e06ddebb7f2a45e76634c65d197b8 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 21 May 2026 18:59:28 +0200 Subject: [PATCH] fix(per-horizon-cfc): label-variance anchor + block-diagonal heads grad mask MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two follow-up fixes for Smoke 2 readiness, atomic per feedback_no_partial_refactor: FIX 1 — Task 11 target_jitter_k anchor (spec §3.3 compliance): - Previously: target_jitter_k = first non-zero raw_per_h[k] (model-state-anchored) - Now: target_jitter_k = sqrt(realised_label_variance_k) (label-distribution-anchored) - Per-horizon label variance computed host-side from existing per-horizon labels in stg_labels (layout [K, B, N_HORIZONS] row-major; typical K=32, B=1 -> 32 floats per horizon -> trivial host reduction). - Wiener-α EMA with α = 0.4 floor + first-observation bootstrap per pearl_wiener_alpha_floor_for_nonstationary + pearl_first_observation_bootstrap. FIX 2 — Task 10 Option B closure (block-diagonal heads via grad mask): - Two new kernels: heads_w_skip_mask_init_kernel (one-shot at transition, builds [N_HORIZONS × HIDDEN_DIM] = 640-float mask + zeros off-bucket heads_w_skip in place) and heads_w_skip_grad_mask_apply_kernel (per-step Phase 2, multiplies grad_heads_w_skip by mask elementwise before opt_heads_w_skip.step). - Achieves block-diagonal heads_w_skip behavior WITHOUT touching the existing GRN kernel: off-bucket positions stay 0 throughout training because gradient is masked out, so Adam never updates them. - Existing GRN dispatch consumes heads_w_skip_d (full 640 floats) as today; off-bucket entries are always 0, so the dispatch is mathematically equivalent to a compact-only read. Together these fixes restore the spec-mandated per-horizon differentiation chain across all 3 mechanisms (CfC.tau bucketing + block-diagonal heads + per-bucket LR with label-variance anchor) before Smoke 2 fires the mean_run_len ratio gate. GPU oracle tests: 13 total (11 + heads_w_skip_mask_init + heads_w_skip_grad_mask_apply). Co-Authored-By: Claude Opus 4.7 (1M context) --- .../cuda/bucket_transition_kernels.cu | 57 +++++ crates/ml-alpha/src/trainer/perception.rs | 207 +++++++++++++++--- .../tests/bucket_transition_kernels.rs | 110 ++++++++++ 3 files changed, 347 insertions(+), 27 deletions(-) diff --git a/crates/ml-alpha/cuda/bucket_transition_kernels.cu b/crates/ml-alpha/cuda/bucket_transition_kernels.cu index c86779330..5a2dea326 100644 --- a/crates/ml-alpha/cuda/bucket_transition_kernels.cu +++ b/crates/ml-alpha/cuda/bucket_transition_kernels.cu @@ -364,6 +364,63 @@ extern "C" __global__ void h_mag_per_bucket_kernel( } } +// ───────────────────────────────────────────────────────────────────── +// heads_w_skip_mask_init_kernel — Phase 1→2 transition one-shot. +// +// Per spec §3.3 (block-diagonal heads) and follow-up to Task 10 Option B: +// build the per-channel routing mask AND zero-out off-bucket entries of +// `heads_w_skip` in place so that subsequent forward/backward never sees +// non-zero cross-bucket weights. Combined with the grad-mask apply kernel +// below, this is mathematically equivalent to a compact-only read of +// `heads_w_skip` at zero implementation cost for the GRN kernel. +// +// Layout: `heads_w_skip` is `[N_HORIZONS × HIDDEN_DIM]` row-major +// (horizon-major). `bucket_id_per_channel[c]` records which horizon owns +// column `c` in the block-diagonal layout. Entry `(h, c)` is in-bucket +// iff `bucket_id_per_channel[c] == h`. +// +// Launch: grid = (N_HORIZONS, 1, 1), block = (HIDDEN_DIM, 1, 1). +// One block per horizon × one thread per channel. No reductions, no +// shared memory, no atomicAdd — fully data-parallel. +// ───────────────────────────────────────────────────────────────────── +extern "C" __global__ void heads_w_skip_mask_init_kernel( + const unsigned char* __restrict__ bucket_id_per_channel, // [HIDDEN_DIM] + float* __restrict__ heads_w_skip, // [N_HORIZONS × HIDDEN_DIM] + float* __restrict__ mask // [N_HORIZONS × HIDDEN_DIM] +) { + int h = blockIdx.x; + int c = threadIdx.x; + if (h >= N_HORIZONS || c >= HIDDEN_DIM) return; + int idx = h * HIDDEN_DIM + c; + float m = (bucket_id_per_channel[c] == (unsigned char)h) ? 1.0f : 0.0f; + mask[idx] = m; + heads_w_skip[idx] *= m; // zero off-bucket entries in place +} + +// ───────────────────────────────────────────────────────────────────── +// heads_w_skip_grad_mask_apply_kernel — per-step Phase 2 gradient mask. +// +// Per spec §3.3 + follow-up Option B closure: multiply the post-reduction +// `grad_heads_w_skip` element-wise by the routing mask BEFORE the Adam +// step. Off-bucket positions receive zero gradient, never update, and +// (since they started at zero after the init kernel) remain at zero +// throughout training. This delivers block-diagonal `heads_w_skip` +// behaviour without touching the GRN forward/backward kernels. +// +// Launch: grid = (N_HORIZONS, 1, 1), block = (HIDDEN_DIM, 1, 1). Same +// layout as `heads_w_skip_mask_init_kernel`. +// ───────────────────────────────────────────────────────────────────── +extern "C" __global__ void heads_w_skip_grad_mask_apply_kernel( + const float* __restrict__ mask, // [N_HORIZONS × HIDDEN_DIM] + float* __restrict__ grad_heads_w_skip // [N_HORIZONS × HIDDEN_DIM] +) { + int h = blockIdx.x; + int c = threadIdx.x; + if (h >= N_HORIZONS || c >= HIDDEN_DIM) return; + int idx = h * HIDDEN_DIM + c; + grad_heads_w_skip[idx] *= mask[idx]; +} + // ───────────────────────────────────────────────────────────────────── // bucket_iqr_double_widening_kernel — Controller D Recovery Attempt 1. // diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index 1fa2fd4f0..289a1d4c3 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -760,6 +760,25 @@ pub struct PerceptionTrainer { /// rescaling. DtoD-copied each Phase 2 step BEFORE `opt_heads_w_skip.step`. /// Shape matches `heads_w_skip_d` (N_HORIZONS × HIDDEN_DIM = 640 floats). heads_w_skip_pre_adam_d: CudaSlice, + /// Per-channel routing mask for `heads_w_skip` (block-diagonal heads + /// via gradient masking — follow-up to Task 10 Option B). Built once + /// at Phase 1→2 transition by `heads_w_skip_mask_init_kernel`: + /// `mask[h, c] = 1.0` iff `bucket_id_per_channel[c] == h`, else 0.0. + /// At the same launch the kernel zero-multiplies off-bucket entries + /// of `heads_w_skip_d` so the GRN forward sees only block-diagonal + /// weights. Each Phase 2 backward step + /// `heads_w_skip_grad_mask_apply_kernel` multiplies + /// `grad_heads_w_skip_d` by this mask before Adam — off-bucket + /// entries stay at 0 throughout training. Shape `[N_HORIZONS × HIDDEN_DIM]`. + heads_w_skip_mask_d: CudaSlice, + /// Cached function handle for `heads_w_skip_mask_init_kernel` — + /// one-shot at Phase 1→2 transition (builds mask + zeros off-bucket + /// entries of `heads_w_skip_d` in place). + heads_w_skip_mask_init_fn: CudaFunction, + /// Cached function handle for `heads_w_skip_grad_mask_apply_kernel` — + /// per-step Phase 2 (multiplies `grad_heads_w_skip_d` by the routing + /// mask before Adam, enforcing block-diagonal updates). + heads_w_skip_grad_mask_apply_fn: CudaFunction, /// Per-horizon LR multiplier vector consumed by /// `heads_lr_multiplier_scale_kernel`. Mapped-pinned per /// `feedback_no_htod_htoh_only_mapped_pinned`: host computes @@ -774,12 +793,24 @@ pub struct PerceptionTrainer { /// smoothness ratio. smoothness_jitter_ema_host_d: MappedF32Buffer, /// Per-horizon target jitter for Controller C (`target_jitter_k` in - /// spec §3.3). Sentinel-bootstrapped per - /// `pearl_first_observation_bootstrap`: stays at 0 until the first - /// non-zero `raw_per_h` observation in Phase 2; then is written - /// directly with that value. After bootstrap, this acts as the - /// `realised_label_variance_k` proxy. + /// spec §3.3). Each Phase 2 step the host writes + /// `target_jitter_per_h[k] = sqrt(label_var_ema[k])` — anchored to the + /// observed label distribution per spec §3.3 + `feedback_isv_for_adaptive_bounds` + /// (data property, not model state). See `label_var_ema` below for + /// the variance EMA itself. target_jitter_per_h: [f32; N_HORIZONS], + /// Per-horizon EMA of `Var(label_k)` over each batch's K×B labels. + /// Wiener-α with α = 0.4 floor per + /// `pearl_wiener_alpha_floor_for_nonstationary` (the closed-loop + /// gradient updates make Controller C a non-stationary target); + /// first-observation bootstrap per `pearl_first_observation_bootstrap` + /// (replace, don't blend, on the first sample). Source for + /// `target_jitter_per_h[k]` via `sqrt(label_var_ema[k])`. + label_var_ema: [f32; N_HORIZONS], + /// Has `label_var_ema` received its first observation? Gated separately + /// from `target_jitter_per_h` because the variance EMA is what's + /// bootstrapped — `target_jitter_per_h` is now derived each step. + label_var_ema_bootstrapped: bool, // ── Controller D state (dead-bucket detector; spec §3.4) ── // @@ -1019,6 +1050,12 @@ impl PerceptionTrainer { let bucket_iqr_double_widening_fn = bucket_transition_module .load_function("bucket_iqr_double_widening_kernel") .context("load bucket_iqr_double_widening_kernel")?; + let heads_w_skip_mask_init_fn = bucket_transition_module + .load_function("heads_w_skip_mask_init_kernel") + .context("load heads_w_skip_mask_init_kernel")?; + let heads_w_skip_grad_mask_apply_fn = bucket_transition_module + .load_function("heads_w_skip_grad_mask_apply_kernel") + .context("load heads_w_skip_grad_mask_apply_kernel")?; // Mamba2 stacks live on the trunk from new_random; we wire optimizer // + training scratches against the trunk-owned blocks. @@ -1401,6 +1438,13 @@ impl PerceptionTrainer { let heads_w_skip_pre_adam_d = stream .alloc_zeros::(N_HORIZONS * crate::cfc::bucket_routing::HIDDEN_DIM) .context("heads_w_skip_pre_adam_d alloc")?; + // Block-diagonal heads grad-mask (follow-up to Task 10 Option B). + // Zero-allocated; populated at Phase 1→2 transition by + // `heads_w_skip_mask_init_kernel`. Phase 1 reads it never — the + // grad-mask apply kernel only fires in Phase 2. + let heads_w_skip_mask_d = stream + .alloc_zeros::(N_HORIZONS * crate::cfc::bucket_routing::HIDDEN_DIM) + .context("heads_w_skip_mask_d alloc")?; // Per-horizon LR multiplier staging. Mapped-pinned so the host // writes from `target_jitter_per_h` + jitter EMA each Phase 2 step // and the kernel reads the device-visible pointer without DtoH. @@ -1796,9 +1840,14 @@ impl PerceptionTrainer { bucket_routing_metadata: None, // Controller C state (Task 11). heads_w_skip_pre_adam_d, + heads_w_skip_mask_d, + heads_w_skip_mask_init_fn, + heads_w_skip_grad_mask_apply_fn, lr_mult_per_horizon_staging, smoothness_jitter_ema_host_d, - target_jitter_per_h: [0.0; N_HORIZONS], // sentinel; first obs bootstraps + target_jitter_per_h: [0.0; N_HORIZONS], + label_var_ema: [0.0; N_HORIZONS], + label_var_ema_bootstrapped: false, // Controller D state (Task 12). controller_d: ControllerDState::new(), phase2_step_count: 0, @@ -2152,6 +2201,41 @@ impl PerceptionTrainer { // Controllers B / C / D. self.bucket_routing_metadata = Some(metadata); + // 5.5. Block-diagonal heads grad-mask init (follow-up to + // Task 10 Option B). One-shot: build the routing + // mask AND zero-multiply off-bucket entries of + // `heads_w_skip_d` in place. From here on, the per- + // step grad-mask apply (below `opt_heads_w_skip.step`) + // keeps off-bucket positions at 0 throughout training, + // mathematically equivalent to a compact-only read of + // `heads_w_skip` at zero implementation cost for the + // GRN kernel. Spec §3.3 + `feedback_isv_for_adaptive_bounds`: + // the mask is derived from the same + // `bucket_id_per_channel_d` that Controllers B/C use. + { + let metadata_ref = self + .bucket_routing_metadata + .as_ref() + .expect("metadata just persisted above"); + let cfg_mask_init = LaunchConfig { + grid_dim: (N_HORIZONS as u32, 1, 1), + block_dim: (HIDDEN_DIM as u32, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = self + .stream + .launch_builder(&self.heads_w_skip_mask_init_fn); + launch + .arg(&metadata_ref.bucket_id_per_channel_d) + .arg(&mut self.trunk.heads_w_skip_d) + .arg(&mut self.heads_w_skip_mask_d); + unsafe { + launch + .launch(cfg_mask_init) + .context("heads_w_skip_mask_init_kernel launch (Phase 1→2 block-diagonal heads)")?; + } + } + // 6. Latch trainer into Phase 2. // // NOTE: Task 9 leaves the per-step dispatch path @@ -2224,40 +2308,83 @@ impl PerceptionTrainer { // ── 4. Controller C host-side LR multiplier update (spec §3.3) ── // // Now that the end-of-step sync has propagated the captured DtoD - // shadows, both `smoothness_jitter_ema_host_d` and - // `smoothness_loss_per_horizon_host_d` carry this step's per-horizon - // values. The host: - // 1. Sentinel-bootstraps `target_jitter_per_h[k]` from the FIRST - // non-zero raw smoothness loss in Phase 2 (per - // `pearl_first_observation_bootstrap`). - // 2. Computes `smoothness_ratio_k = max(0, jitter_ema - target) / + // shadows, `smoothness_jitter_ema_host_d` carries this step's + // per-horizon jitter EMA. Per spec §3.3 the host: + // 1. Computes per-horizon `label_var_k = E[label²] - E[label]²` + // over this batch's `[K × B]` labels for horizon k (host-side + // reduction on `stg_labels`; layout is `[K, B, N_HORIZONS]` + // row-major — see the staging fill earlier in this function). + // 2. Updates `label_var_ema[k]` via Wiener-α EMA (α = 0.4 floor + // per `pearl_wiener_alpha_floor_for_nonstationary`; first- + // observation bootstrap per `pearl_first_observation_bootstrap`). + // 3. Writes `target_jitter_per_h[k] = sqrt(label_var_ema[k])` — + // anchored to the label DISTRIBUTION per spec §3.3 + + // `feedback_isv_for_adaptive_bounds` (data property, not + // model state). Replaces the prior model-state-anchored + // bootstrap from `raw_per_h` which deviated from the spec. + // 4. Computes `smoothness_ratio_k = max(0, jitter_ema - target) / // max(target, ε_floor)` per `pearl_blend_formulas_must_have_permanent_floor`. - // 3. Writes `lr_mult_k = 1 / (1 + smoothness_ratio_k)` (bounded + // 5. Writes `lr_mult_k = 1 / (1 + smoothness_ratio_k)` (bounded // (0, 1]) to the mapped-pinned staging slice. The next // step's captured-graph rescale kernel reads the updated // values via the device-visible pointer (mapped-pinned // coherence — no DtoH). // // Phase 1: no-op (the rescale kernel is also gated to Phase 2, - // and the staging keeps its zero-init / bootstrap value). + // and the staging keeps its 1.0-init value). if self.phase == TrainingPhase::Phase2Routed { + // (1) Host-side per-horizon label variance over this batch. + // + // `stg_labels` layout is `[K × B × N_HORIZONS]` row-major + // (horizon innermost): index `(k * B + b) * N_HORIZONS + h`. + // For typical configs (K=32, B=1) this is 32 elements per + // horizon — a trivial host reduction. + let labels = self.stg_labels.read_all(); + let n_per_h: usize = k_seq * b_sz; + debug_assert_eq!(labels.len(), n_per_h * N_HORIZONS); + let mut label_var = [0.0_f32; N_HORIZONS]; + if n_per_h > 0 { + for h in 0..N_HORIZONS { + let mut s = 0.0_f32; + let mut s2 = 0.0_f32; + for kb in 0..n_per_h { + let v = labels[kb * N_HORIZONS + h]; + s += v; + s2 += v * v; + } + let n_f = n_per_h as f32; + let mean = s / n_f; + let mean_sq = s2 / n_f; + // E[x²] − E[x]²; clamped to ≥ 0 for floating-point safety. + label_var[h] = (mean_sq - mean * mean).max(0.0); + } + } + // (2) Wiener-α EMA with α = 0.4 floor + first-observation bootstrap. + const ALPHA_LABEL_VAR: f32 = 0.4; + if !self.label_var_ema_bootstrapped { + self.label_var_ema = label_var; + self.label_var_ema_bootstrapped = true; + } else { + for h in 0..N_HORIZONS { + self.label_var_ema[h] = (1.0 - ALPHA_LABEL_VAR) + * self.label_var_ema[h] + + ALPHA_LABEL_VAR * label_var[h]; + } + } + // (3) target_jitter_per_h[k] = sqrt(label_var_ema[k]) — spec §3.3. + for k in 0..N_HORIZONS { + self.target_jitter_per_h[k] = self.label_var_ema[k].sqrt(); + } + let jitter_ema = self.smoothness_jitter_ema_host_d.read_all(); - let raw_per_h = self.smoothness_loss_per_horizon_host_d.read_all(); let mut lr_mult = [1.0_f32; N_HORIZONS]; for k in 0..N_HORIZONS { - // Sentinel-bootstrap target_jitter_per_h from first non-zero - // observation. Both EMA and raw can legitimately read as - // sentinel-zero on the first Phase 2 step before the - // smoothness controller has run; once the first non-zero - // sample arrives we replace directly per the pearl. - if self.target_jitter_per_h[k] == 0.0 && raw_per_h[k] > 0.0 { - self.target_jitter_per_h[k] = raw_per_h[k]; - } let target = self.target_jitter_per_h[k]; if target <= 0.0 { - // Still pre-bootstrap: leave lr_mult at 1.0 (no - // attenuation). Once the first observation lands the - // next step picks up a non-zero target. + // Pre-bootstrap edge case (label_var_ema still 0 at + // Phase 2 step 0 corner). Leave lr_mult at 1.0 — no + // attenuation, no division by sentinel zero. The next + // step's batch supplies a non-zero label variance. continue; } // ε_floor = target / 16 per `pearl_blend_formulas_must_have_permanent_floor` @@ -3691,6 +3818,32 @@ impl PerceptionTrainer { self.opt_heads_b_gate.step(&mut self.trunk.heads_b_gate_d, &self.grad_heads_b_gate_d)?; self.opt_heads_w_main.step(&mut self.trunk.heads_w_main_d, &self.grad_heads_w_main_d)?; self.opt_heads_b_main.step(&mut self.trunk.heads_b_main_d, &self.grad_heads_b_main_d)?; + // ── Block-diagonal heads: grad mask apply (Phase 2 only) ── + // Follow-up to Task 10 Option B. Multiplies `grad_heads_w_skip_d` + // element-wise by the routing mask built at transition. Off-bucket + // entries see zero gradient → Adam never updates them → they stay + // at 0 forever (init zero-multiplied them at the transition). This + // realizes block-diagonal `heads_w_skip` without modifying the GRN + // forward/backward kernels. MUST run BEFORE + // `opt_heads_w_skip.step` so the masked grad is the one Adam reads. + if self.phase == TrainingPhase::Phase2Routed { + let cfg_grad_mask = LaunchConfig { + grid_dim: (N_HORIZONS as u32, 1, 1), + block_dim: (HIDDEN_DIM as u32, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = self + .stream + .launch_builder(&self.heads_w_skip_grad_mask_apply_fn); + launch + .arg(&self.heads_w_skip_mask_d) + .arg(&mut self.grad_heads_w_skip_d); + unsafe { + launch + .launch(cfg_grad_mask) + .context("heads_w_skip_grad_mask_apply_kernel launch (block-diagonal heads)")?; + } + } // ── Controller C: pre-Adam snapshot of heads_w_skip_d (spec §3.3) ── // Captured DtoD copy: heads_w_skip_d → heads_w_skip_pre_adam_d. Used // by the post-Adam rescale below to compute the parameter delta diff --git a/crates/ml-alpha/tests/bucket_transition_kernels.rs b/crates/ml-alpha/tests/bucket_transition_kernels.rs index 86603b929..bfb64366e 100644 --- a/crates/ml-alpha/tests/bucket_transition_kernels.rs +++ b/crates/ml-alpha/tests/bucket_transition_kernels.rs @@ -460,6 +460,116 @@ fn bucket_iqr_double_widening_kernel_doubles_one_bucket_iqr() -> Result<()> { Ok(()) } +#[test] +#[ignore = "requires CUDA"] +fn heads_w_skip_mask_init_zeros_off_bucket() -> Result<()> { + let dev = MlDevice::cuda(0)?; + let stream = dev.cuda_stream()?.clone(); + let func = load_kernel(&stream, "heads_w_skip_mask_init_kernel")?; + + // Bucket assignment: channel c → bucket c/25 capped at 4 (quintile + // boundaries used elsewhere in this file: bucket_dim_k = [25, 25, 25, + // 25, 28]). heads_w_skip starts all 1.0; mask + zero-out result is + // 1.0 at in-bucket positions, 0.0 elsewhere. + let bucket_id_host: Vec = + (0..HIDDEN_DIM).map(|c| ((c / 25).min(4)) as u8).collect(); + let heads_w_skip_host: Vec = vec![1.0; N_HORIZONS * HIDDEN_DIM]; + + let bucket_id_d = upload(&stream, &bucket_id_host)?; + let mut heads_w_skip_d = upload(&stream, &heads_w_skip_host)?; + let mut mask_d = stream.alloc_zeros::(N_HORIZONS * HIDDEN_DIM)?; + + let cfg = LaunchConfig { + grid_dim: (N_HORIZONS as u32, 1, 1), + block_dim: (HIDDEN_DIM as u32, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = stream.launch_builder(&func); + launch + .arg(&bucket_id_d) + .arg(&mut heads_w_skip_d) + .arg(&mut mask_d); + unsafe { launch.launch(cfg)?; } + stream.synchronize()?; + + let mask = download(&stream, &mask_d)?; + let heads_w_skip_out = download(&stream, &heads_w_skip_d)?; + for h in 0..N_HORIZONS { + for c in 0..HIDDEN_DIM { + let idx = h * HIDDEN_DIM + c; + let in_bucket = (c / 25).min(4) == h; + let expected = if in_bucket { 1.0 } else { 0.0 }; + assert!( + (mask[idx] - expected).abs() < 1e-6, + "mask[{},{}]={} expected {}", + h, + c, + mask[idx], + expected + ); + assert!( + (heads_w_skip_out[idx] - expected).abs() < 1e-6, + "heads_w_skip[{},{}]={} expected {}", + h, + c, + heads_w_skip_out[idx], + expected + ); + } + } + Ok(()) +} + +#[test] +#[ignore = "requires CUDA"] +fn heads_w_skip_grad_mask_apply_zeros_off_bucket_grad() -> Result<()> { + let dev = MlDevice::cuda(0)?; + let stream = dev.cuda_stream()?.clone(); + let func = load_kernel(&stream, "heads_w_skip_grad_mask_apply_kernel")?; + + // Build a mask manually for the same quintile layout used above; feed + // an all-ones grad tensor and expect off-bucket entries to be zeroed. + let mask_host: Vec = (0..N_HORIZONS) + .flat_map(|h| { + (0..HIDDEN_DIM).map(move |c| { + if (c / 25).min(4) == h { 1.0_f32 } else { 0.0_f32 } + }) + }) + .collect(); + let grad_host: Vec = vec![1.0; N_HORIZONS * HIDDEN_DIM]; + + let mask_d = upload(&stream, &mask_host)?; + let mut grad_d = upload(&stream, &grad_host)?; + + let cfg = LaunchConfig { + grid_dim: (N_HORIZONS as u32, 1, 1), + block_dim: (HIDDEN_DIM as u32, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = stream.launch_builder(&func); + launch.arg(&mask_d).arg(&mut grad_d); + unsafe { launch.launch(cfg)?; } + stream.synchronize()?; + + let grad_out = download(&stream, &grad_d)?; + for h in 0..N_HORIZONS { + for c in 0..HIDDEN_DIM { + let idx = h * HIDDEN_DIM + c; + let in_bucket = (c / 25).min(4) == h; + let expected = if in_bucket { 1.0 } else { 0.0 }; + assert!( + (grad_out[idx] - expected).abs() < 1e-6, + "grad[{},{}]={} expected {}", + h, + c, + grad_out[idx], + expected + ); + } + } + Ok(()) +} + #[test] #[ignore = "requires CUDA"] fn slack_factor_apply_kernel_widens_iqr_geometrically() -> Result<()> {