gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Revert "fix(per-horizon-cfc): extend block-diagonal mask to heads_w1 (Smoke 2 fix)"

Browse Source 
This reverts commit 4c1ed1d953.
This commit is contained in:
jgrusewski
2026-05-21 23:48:15 +02:00
parent 4c1ed1d953
commit e18325aaf1
3 changed files with 0 additions and 506 deletions
Download Patch File Download Diff File Expand all files Collapse all files

104
crates/ml-alpha/cuda/bucket_transition_kernels.cu
View File

@@ -5,7 +5,6 @@
// boundaries for HIDDEN_DIM=128: [0, 25, 50, 75, 100, 128].
#define HIDDEN_DIM 128
#define HEAD_MID_DIM 64
#define N_HORIZONS 5
#define BUCKET_DIM_LAST 28 // last bucket absorbs HIDDEN_DIM - 4*25 = 28
#define MAX_BUCKET_DIM 28
@@ -533,106 +532,3 @@ extern "C" __global__ void bucket_iqr_double_widening_kernel(
}
}
}
// ─────────────────────────────────────────────────────────────────────
// heads_w1_mask_init_kernel — Phase 1→2 transition one-shot for
// `heads_w1` (GRN input projection W1: HIDDEN → HEAD_MID per horizon).
//
// Diagnostic smoke 2 (`lob-backtest-sweep-kw7f6` at d2d34b6d4) confirmed
// that with only `heads_w_skip` block-diagonalised, `heads_w1` remained
// 100% dense at off-bucket positions (32,768/32,768 nonzero) and the
// per-horizon logit distributions collapsed to ~uniform (mean spread
// 0.05). `heads_w1` is the SOLE remaining HIDDEN_DIM-reading path in the
// GRN forward (`multi_horizon_heads_grn_fwd_batched` Pass 1: line 497 of
// `multi_horizon_heads.cu`) that bypasses bucket routing — every other
// HIDDEN_DIM-reading weight (`heads_w_skip`) is already restricted, and
// `heads_w2`, `heads_w_gate`, `heads_w_main` only see HEAD_MID-dim
// tensors. Mirrors the `heads_w_skip` mask init: build the routing mask
// AND zero-multiply off-bucket entries of `heads_w1` in place.
//
// Layout: `heads_w1` is `[N_HORIZONS × HEAD_MID_DIM × HIDDEN_DIM]`
// row-major (`heads_w1[h, m, c]` flat-index = h*HEAD_MID_DIM*HIDDEN_DIM
// + m*HIDDEN_DIM + c). The mask depends only on (h, c) — the m-axis is
// pure replication. Position `(h, m, c)` is in-bucket iff
// `bucket_id_per_channel[c] == h`.
//
// Launch: grid = (N_HORIZONS, HEAD_MID_DIM, 1), block = (HIDDEN_DIM, 1, 1).
// One block per (horizon, mid-row) × one thread per channel. No
// reductions, no shared memory, no atomicAdd — fully data-parallel.
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void heads_w1_mask_init_kernel(
const unsigned char* __restrict__ bucket_id_per_channel, // [HIDDEN_DIM]
float* __restrict__ heads_w1, // [N_HORIZONS × HEAD_MID_DIM × HIDDEN_DIM]
float* __restrict__ mask // [N_HORIZONS × HEAD_MID_DIM × HIDDEN_DIM]
) {
int h = blockIdx.x;
int m = blockIdx.y;
int c = threadIdx.x;
if (h >= N_HORIZONS || m >= HEAD_MID_DIM || c >= HIDDEN_DIM) return;
long long idx = ((long long)h * HEAD_MID_DIM + m) * HIDDEN_DIM + c;
float msk = (bucket_id_per_channel[c] == (unsigned char)h) ? 1.0f : 0.0f;
mask[idx] = msk;
heads_w1[idx] *= msk; // zero off-bucket entries in place
}
// ─────────────────────────────────────────────────────────────────────
// heads_w1_grad_mask_apply_kernel — per-step Phase 2 gradient mask for
// `heads_w1`.
//
// Mirrors `heads_w_skip_grad_mask_apply_kernel`: multiply the post-
// reduction `grad_heads_w1` 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_w1` without
// touching the GRN forward/backward kernels.
//
// Launch: grid = (N_HORIZONS, HEAD_MID_DIM, 1), block = (HIDDEN_DIM, 1, 1).
// Same layout as `heads_w1_mask_init_kernel`.
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void heads_w1_grad_mask_apply_kernel(
const float* __restrict__ mask, // [N_HORIZONS × HEAD_MID_DIM × HIDDEN_DIM]
float* __restrict__ grad_heads_w1 // [N_HORIZONS × HEAD_MID_DIM × HIDDEN_DIM]
) {
int h = blockIdx.x;
int m = blockIdx.y;
int c = threadIdx.x;
if (h >= N_HORIZONS || m >= HEAD_MID_DIM || c >= HIDDEN_DIM) return;
long long idx = ((long long)h * HEAD_MID_DIM + m) * HIDDEN_DIM + c;
grad_heads_w1[idx] *= mask[idx];
}
// ─────────────────────────────────────────────────────────────────────
// heads_w1_zero_off_bucket_kernel — block-diagonal projection on
// `[N_HORIZONS × HEAD_MID_DIM × HIDDEN_DIM]` tensors (params + Adam
// moments).
//
// Same role as `zero_off_bucket_kernel` but for the 3-D heads_w1 shape.
// Launched on:
// 1. `heads_w1_d` at transition (one-shot, zeros off-bucket weights;
// redundant with `heads_w1_mask_init_kernel` but explicit for
// traceability)
// 2. `heads_w1` Adam `m` and `v` buffers at transition (one-shot; zeros
// pre-transition momentum so post-transition Adam steps cannot
// revive off-bucket positions)
// 3. `heads_w1_d` after every Phase 2 Adam step (belt-and-suspenders
// against ε-driven drift from zero moments)
//
// Layout matches `heads_w1_mask_init_kernel`: `tensor[h, m, c]` flat-
// index = `h*HEAD_MID_DIM*HIDDEN_DIM + m*HIDDEN_DIM + c`. Off-bucket iff
// `bucket_id_per_channel[c] != h`.
//
// Launch: grid = (N_HORIZONS, HEAD_MID_DIM, 1), block = (HIDDEN_DIM, 1, 1).
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void heads_w1_zero_off_bucket_kernel(
const unsigned char* __restrict__ bucket_id_per_channel, // [HIDDEN_DIM]
float* __restrict__ tensor // [N_HORIZONS × HEAD_MID_DIM × HIDDEN_DIM]
) {
int h = blockIdx.x;
int m = blockIdx.y;
int c = threadIdx.x;
if (h >= N_HORIZONS || m >= HEAD_MID_DIM || c >= HIDDEN_DIM) return;
long long idx = ((long long)h * HEAD_MID_DIM + m) * HIDDEN_DIM + c;
if (bucket_id_per_channel[c] != (unsigned char)h) {
tensor[idx] = 0.0f;
}
}

192
crates/ml-alpha/src/trainer/perception.rs
View File

@@ -792,31 +792,6 @@ pub struct PerceptionTrainer {
/// exactly zero even under Adam's ε-driven step from zero
/// moments).
zero_off_bucket_fn: CudaFunction,
/// Per-channel routing mask for `heads_w1` (GRN-wide block-diagonal
/// extension; Smoke 2 fix 2026-05-21). Built once at Phase 1→2
/// transition by `heads_w1_mask_init_kernel`:
/// `mask[h, m, 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_w1_d` so the GRN forward sees only block-diagonal weights on
/// the HIDDEN-reading W1 path (Pass 1 of `multi_horizon_heads_grn_fwd_
/// batched`). Each Phase 2 backward step
/// `heads_w1_grad_mask_apply_kernel` multiplies `grad_heads_w1_d` by
/// this mask before Adam — off-bucket entries stay at 0 throughout
/// training. Shape `[N_HORIZONS × HEAD_MID_DIM × HIDDEN_DIM]`.
heads_w1_mask_d: CudaSlice<f32>,
/// Cached function handle for `heads_w1_mask_init_kernel` — one-shot
/// at Phase 1→2 transition (builds mask + zeros off-bucket entries of
/// `heads_w1_d` in place).
heads_w1_mask_init_fn: CudaFunction,
/// Cached function handle for `heads_w1_grad_mask_apply_kernel` —
/// per-step Phase 2 (multiplies `grad_heads_w1_d` by the routing mask
/// before Adam, enforcing block-diagonal updates on the W1 path).
heads_w1_grad_mask_apply_fn: CudaFunction,
/// Cached function handle for `heads_w1_zero_off_bucket_kernel` —
/// 3-D analogue of `zero_off_bucket_kernel`, used on
/// `[N_HORIZONS × HEAD_MID_DIM × HIDDEN_DIM]` tensors (heads_w1
/// params + Adam m/v moments) at transition and post-Adam in Phase 2.
heads_w1_zero_off_bucket_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
@@ -1097,15 +1072,6 @@ impl PerceptionTrainer {
let zero_off_bucket_fn = bucket_transition_module
.load_function("zero_off_bucket_kernel")
.context("load zero_off_bucket_kernel")?;
let heads_w1_mask_init_fn = bucket_transition_module
.load_function("heads_w1_mask_init_kernel")
.context("load heads_w1_mask_init_kernel")?;
let heads_w1_grad_mask_apply_fn = bucket_transition_module
.load_function("heads_w1_grad_mask_apply_kernel")
.context("load heads_w1_grad_mask_apply_kernel")?;
let heads_w1_zero_off_bucket_fn = bucket_transition_module
.load_function("heads_w1_zero_off_bucket_kernel")
.context("load heads_w1_zero_off_bucket_kernel")?;
// Mamba2 stacks live on the trunk from new_random; we wire optimizer
// + training scratches against the trunk-owned blocks.
@@ -1495,16 +1461,6 @@ impl PerceptionTrainer {
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")?;
// GRN-wide block-diagonal: heads_w1 mask (Smoke 2 fix 2026-05-21).
// Same lifecycle as `heads_w_skip_mask_d`: zero-allocated here,
// populated at Phase 1→2 transition by `heads_w1_mask_init_kernel`,
// read each Phase 2 step by `heads_w1_grad_mask_apply_kernel`.
// Shape `[N_HORIZONS × HEAD_MID_DIM × HIDDEN_DIM]` = 40,960 floats.
let heads_w1_mask_d = stream
.alloc_zeros::<f32>(
N_HORIZONS * HEAD_MID_DIM * crate::cfc::bucket_routing::HIDDEN_DIM,
)
.context("heads_w1_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.
@@ -1904,10 +1860,6 @@ impl PerceptionTrainer {
heads_w_skip_mask_init_fn,
heads_w_skip_grad_mask_apply_fn,
zero_off_bucket_fn,
heads_w1_mask_d,
heads_w1_mask_init_fn,
heads_w1_grad_mask_apply_fn,
heads_w1_zero_off_bucket_fn,
lr_mult_per_horizon_staging,
smoothness_jitter_ema_host_d,
target_jitter_per_h: [0.0; N_HORIZONS],
@@ -2433,94 +2385,6 @@ impl PerceptionTrainer {
}
}
// 7c. GRN-wide block-diagonal extension (Smoke 2 fix
// 2026-05-21). Diagnostic
// `lob-backtest-sweep-kw7f6` at d2d34b6d4 confirmed
// `heads_w1` is the sole remaining HIDDEN_DIM-reading
// path in the GRN forward (`multi_horizon_heads_grn_
// fwd_batched` Pass 1) that bypasses bucket routing,
// and that it stayed 100% dense (32,768/32,768 nonzero
// at off-bucket positions) when only `heads_w_skip`
// was restricted — collapsing per-horizon logits to
// ~uniform (mean spread 0.05, mean_run_len ratio
// 1.04×). Mirror the heads_w_skip defense exactly:
// mask init (zero off-bucket params + build mask),
// Adam moment projection (m, v), grad mask per step,
// post-Adam projection per step (belt+suspenders).
{
let metadata_ref = self
.bucket_routing_metadata
.as_ref()
.expect("metadata persisted in step 6");
let cfg_mask_init_w1 = LaunchConfig {
grid_dim: (N_HORIZONS as u32, HEAD_MID_DIM as u32, 1),
block_dim: (HIDDEN_DIM as u32, 1, 1),
shared_mem_bytes: 0,
};
// (i) Build mask + zero off-bucket params of heads_w1
// in place.
{
let mut launch = self
.stream
.launch_builder(&self.heads_w1_mask_init_fn);
launch
.arg(&metadata_ref.bucket_id_per_channel_d)
.arg(&mut self.trunk.heads_w1_d)
.arg(&mut self.heads_w1_mask_d);
unsafe {
launch
.launch(cfg_mask_init_w1)
.context("heads_w1_mask_init_kernel launch (Phase 1→2 block-diagonal heads_w1)")?;
}
}
// (ii) Traceability re-zero of heads_w1_d (redundant
// with the mask_init zero above; same role as the
// heads_w_skip step 7b.(a) launch).
{
let mut launch = self
.stream
.launch_builder(&self.heads_w1_zero_off_bucket_fn);
launch
.arg(&metadata_ref.bucket_id_per_channel_d)
.arg(&mut self.trunk.heads_w1_d);
unsafe {
launch
.launch(cfg_mask_init_w1)
.context("heads_w1_zero_off_bucket_kernel launch (heads_w1_d at transition)")?;
}
}
// (iii) Project Adam's first-moment buffer `m` for
// heads_w1.
{
let mut launch = self
.stream
.launch_builder(&self.heads_w1_zero_off_bucket_fn);
launch
.arg(&metadata_ref.bucket_id_per_channel_d)
.arg(self.opt_heads_w1.m_mut());
unsafe {
launch
.launch(cfg_mask_init_w1)
.context("heads_w1_zero_off_bucket_kernel launch (heads_w1 Adam m at transition)")?;
}
}
// (iv) Project Adam's second-moment buffer `v` for
// heads_w1.
{
let mut launch = self
.stream
.launch_builder(&self.heads_w1_zero_off_bucket_fn);
launch
.arg(&metadata_ref.bucket_id_per_channel_d)
.arg(self.opt_heads_w1.v_mut());
unsafe {
launch
.launch(cfg_mask_init_w1)
.context("heads_w1_zero_off_bucket_kernel launch (heads_w1 Adam v at transition)")?;
}
}
}
// 8. Latch trainer into Phase 2.
//
// ALPHA fix 2026-05-21: τ stays in ORIGINAL channel
@@ -4101,63 +3965,7 @@ impl PerceptionTrainer {
}
}
}
// ── Block-diagonal heads_w1: grad mask apply (Phase 2 only) ──
// GRN-wide block-diagonal extension (Smoke 2 fix 2026-05-21).
// Mirrors the heads_w_skip grad mask below: multiplies
// `grad_heads_w1_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). Realises block-diagonal heads_w1 on
// the GRN's Pass-1 HIDDEN-reading path without touching
// `multi_horizon_heads_grn_fwd_batched` itself. MUST run BEFORE
// `opt_heads_w1.step` so the masked grad is the one Adam reads.
if self.phase == TrainingPhase::Phase2Routed {
let cfg_grad_mask_w1 = LaunchConfig {
grid_dim: (N_HORIZONS as u32, HEAD_MID_DIM as u32, 1),
block_dim: (HIDDEN_DIM as u32, 1, 1),
shared_mem_bytes: 0,
};
let mut launch = self
.stream
.launch_builder(&self.heads_w1_grad_mask_apply_fn);
launch
.arg(&self.heads_w1_mask_d)
.arg(&mut self.grad_heads_w1_d);
unsafe {
launch
.launch(cfg_grad_mask_w1)
.context("heads_w1_grad_mask_apply_kernel launch (block-diagonal heads_w1)")?;
}
}
self.opt_heads_w1.step(&mut self.trunk.heads_w1_d, &self.grad_heads_w1_d)?;
// ── Block-diagonal heads_w1: post-Adam projection (Phase 2 only) ──
// Belt-and-suspenders re-zero of off-bucket `heads_w1_d` positions
// after every Phase 2 Adam step. Same role as the heads_w_skip
// post-Adam zero below: even with the grad mask + Adam-moment
// mask at transition, mixed-precision rounding and finite-precision
// moment buffers can drift; we project hard per
// `feedback_no_quickfixes`. Effectively realises
// `heads_w1[h, m, c] = 0 for c not in bucket h`.
if self.phase == TrainingPhase::Phase2Routed {
if let Some(metadata) = self.bucket_routing_metadata.as_ref() {
let cfg_zero_w1 = LaunchConfig {
grid_dim: (N_HORIZONS as u32, HEAD_MID_DIM as u32, 1),
block_dim: (HIDDEN_DIM as u32, 1, 1),
shared_mem_bytes: 0,
};
let mut launch = self
.stream
.launch_builder(&self.heads_w1_zero_off_bucket_fn);
launch
.arg(&metadata.bucket_id_per_channel_d)
.arg(&mut self.trunk.heads_w1_d);
unsafe {
launch
.launch(cfg_zero_w1)
.context("heads_w1_zero_off_bucket_kernel launch (post-Adam block-diagonal projection)")?;
}
}
}
self.opt_heads_b1.step(&mut self.trunk.heads_b1_d, &self.grad_heads_b1_d)?;
self.opt_heads_w2.step(&mut self.trunk.heads_w2_d, &self.grad_heads_w2_d)?;
self.opt_heads_b2.step(&mut self.trunk.heads_b2_d, &self.grad_heads_b2_d)?;

210
crates/ml-alpha/tests/bucket_transition_kernels.rs
View File

@@ -873,213 +873,3 @@ fn slack_factor_apply_kernel_widens_iqr_geometrically() -> Result<()> {
}
Ok(())
}
// ─────────────────────────────────────────────────────────────────────
// heads_w1 block-diagonal kernels — GRN-wide extension (Smoke 2 fix
// 2026-05-21). heads_w1 has shape [N_HORIZONS × HEAD_MID_DIM × HIDDEN_DIM].
// The mask depends only on (h, c) — replicated across the m axis.
// ─────────────────────────────────────────────────────────────────────
const HEAD_MID_DIM: usize = 64;
#[test]
#[ignore = "requires CUDA"]
fn heads_w1_mask_init_zeros_off_bucket() -> Result<()> {
let dev = MlDevice::cuda(0)?;
let stream = dev.cuda_stream()?.clone();
let func = load_kernel(&stream, "heads_w1_mask_init_kernel")?;
// Bucket assignment: channel c → bucket c/25 capped at 4 (same
// quintile layout used by the heads_w_skip mask init test above).
// heads_w1 starts all 1.0 across [N_HORIZONS × HEAD_MID_DIM × HIDDEN_DIM];
// 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 total = N_HORIZONS * HEAD_MID_DIM * HIDDEN_DIM;
let heads_w1_host: Vec<f32> = vec![1.0; total];
let bucket_id_d = upload(&stream, &bucket_id_host)?;
let mut heads_w1_d = upload(&stream, &heads_w1_host)?;
let mut mask_d = stream.alloc_zeros::<f32>(total)?;
let cfg = LaunchConfig {
grid_dim: (N_HORIZONS as u32, HEAD_MID_DIM as u32, 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_w1_d)
.arg(&mut mask_d);
unsafe {
launch.launch(cfg)?;
}
stream.synchronize()?;
let mask = download(&stream, &mask_d)?;
let heads_w1_out = download(&stream, &heads_w1_d)?;
for h in 0..N_HORIZONS {
for m in 0..HEAD_MID_DIM {
for c in 0..HIDDEN_DIM {
let idx = (h * HEAD_MID_DIM + m) * 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, m, c, mask[idx], expected
);
assert!(
(heads_w1_out[idx] - expected).abs() < 1e-6,
"heads_w1[{},{},{}]={} expected {}",
h, m, c, heads_w1_out[idx], expected
);
}
}
}
Ok(())
}
#[test]
#[ignore = "requires CUDA"]
fn heads_w1_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_w1_grad_mask_apply_kernel")?;
// Build a mask manually for the same quintile layout; feed an all-
// ones grad tensor and expect off-bucket entries to be zeroed across
// the full [N_HORIZONS × HEAD_MID_DIM × HIDDEN_DIM] shape.
let total = N_HORIZONS * HEAD_MID_DIM * HIDDEN_DIM;
let mut mask_host: Vec<f32> = vec![0.0; total];
for h in 0..N_HORIZONS {
for m in 0..HEAD_MID_DIM {
for c in 0..HIDDEN_DIM {
if (c / 25).min(4) == h {
mask_host[(h * HEAD_MID_DIM + m) * HIDDEN_DIM + c] = 1.0;
}
}
}
}
let grad_host: Vec<f32> = vec![1.0; total];
let mask_d = upload(&stream, &mask_host)?;
let mut grad_d = upload(&stream, &grad_host)?;
let cfg = LaunchConfig {
grid_dim: (N_HORIZONS as u32, HEAD_MID_DIM as u32, 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 m in 0..HEAD_MID_DIM {
for c in 0..HIDDEN_DIM {
let idx = (h * HEAD_MID_DIM + m) * 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, m, c, grad_out[idx], expected
);
}
}
}
Ok(())
}
#[test]
#[ignore = "requires CUDA"]
fn heads_w1_zero_off_bucket_keeps_off_bucket_zero_after_many_mock_adam_steps() -> Result<()> {
// ALPHA-style invariant test for the 3-D heads_w1 projection. Start
// with heads_w1 = all-ones at in-bucket positions, zero elsewhere
// (post-mask_init state). Run N "Adam steps" of the form
// heads_w1 += eps (uniform tiny push that would breach the
// block-diagonal invariant without projection)
// and verify that after each post-step zero call, off-bucket entries
// return to exactly 0.0 while in-bucket entries accumulate the push.
let dev = MlDevice::cuda(0)?;
let stream = dev.cuda_stream()?.clone();
let func_zero = load_kernel(&stream, "heads_w1_zero_off_bucket_kernel")?;
let func_mask_init = load_kernel(&stream, "heads_w1_mask_init_kernel")?;
let total = N_HORIZONS * HEAD_MID_DIM * HIDDEN_DIM;
let bucket_id_host: Vec<u8> =
(0..HIDDEN_DIM).map(|c| ((c / 25).min(4)) as u8).collect();
let bucket_id_d = upload(&stream, &bucket_id_host)?;
let mut heads_w1_d = upload(&stream, &vec![1.0_f32; total])?;
let mut mask_d = stream.alloc_zeros::<f32>(total)?;
let cfg_w1 = LaunchConfig {
grid_dim: (N_HORIZONS as u32, HEAD_MID_DIM as u32, 1),
block_dim: (HIDDEN_DIM as u32, 1, 1),
shared_mem_bytes: 0,
};
// Initialise: zero off-bucket weights + populate mask.
{
let mut launch = stream.launch_builder(&func_mask_init);
launch
.arg(&bucket_id_d)
.arg(&mut heads_w1_d)
.arg(&mut mask_d);
unsafe {
launch.launch(cfg_w1)?;
}
}
stream.synchronize()?;
// 10 mock steps (keep runtime modest given 40,960-element tensor) of
// uniform-push + projection.
let n_steps = 10;
let eps_push: f32 = 1e-3;
for step in 0..n_steps {
let mut current = download(&stream, &heads_w1_d)?;
for i in 0..current.len() {
current[i] += eps_push;
}
stream.memcpy_htod(&current, &mut heads_w1_d)?;
let mut launch = stream.launch_builder(&func_zero);
launch.arg(&bucket_id_d).arg(&mut heads_w1_d);
unsafe {
launch.launch(cfg_w1)?;
}
stream.synchronize()?;
let post = download(&stream, &heads_w1_d)?;
for h in 0..N_HORIZONS {
// Sample a few m rows to keep test runtime bounded.
for m in [0usize, HEAD_MID_DIM / 2, HEAD_MID_DIM - 1] {
for c in 0..HIDDEN_DIM {
let idx = (h * HEAD_MID_DIM + m) * HIDDEN_DIM + c;
let in_bucket = (c / 25).min(4) == h;
if !in_bucket {
assert_eq!(
post[idx], 0.0,
"off-bucket [{}, {}, {}] = {} after step {} — projection failed",
h, m, c, post[idx], step
);
} else {
// In-bucket positions accumulate the push:
// starts at 1.0; after (step + 1) pushes,
// value = 1.0 + (step+1) * eps_push.
let expected = 1.0_f32 + ((step + 1) as f32) * eps_push;
assert!(
(post[idx] - expected).abs() < 1e-5,
"in-bucket [{}, {}, {}] = {} after step {} expected {}",
h, m, c, post[idx], step, expected
);
}
}
}
}
}
Ok(())
}