fix(per-horizon-cfc): label-variance anchor + block-diagonal heads grad mask
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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.
|
||||
//
|
||||
|
||||
@@ -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<f32>,
|
||||
/// 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<f32>,
|
||||
/// 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::<f32>(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::<f32>(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
|
||||
|
||||
@@ -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<u8> =
|
||||
(0..HIDDEN_DIM).map(|c| ((c / 25).min(4)) as u8).collect();
|
||||
let heads_w_skip_host: Vec<f32> = 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::<f32>(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<f32> = (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<f32> = 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<()> {
|
||||
|
||||
Reference in New Issue
Block a user