foxhunt/crates/ml-alpha/cuda/bucket_transition_kernels.cu

// bucket_transition_kernels.cu — Phase 1→2 transition kernels.
//
// All-on-device per feedback_no_htod_htoh_only_mapped_pinned. No atomicAdd
// per feedback_no_atomicadd; reductions use warp-shuffle. Static tercile
// boundaries for HIDDEN_DIM=128 with N_HORIZONS=3: [0, 43, 86, 128]
// (post-2026-05-22 horizon-rebase to 3 buckets, matches Rust-side
// `crates/ml-alpha/src/cfc/bucket_routing.rs::BUCKET_CHANNEL_OFFSET`).
//
// MAX_BUCKET_DIM=96 matches the Rust-side `MAX_BUCKET_DIM` and is sized
// for the upper bound of ISV-driven τ-assignment concentration (~75% of
// channels into one bucket under extreme τ distributions).

#define HIDDEN_DIM 128
#define N_HORIZONS 3
#define BUCKET_DIM_LAST 42  // last bucket absorbs HIDDEN_DIM - 2*43 = 42
#define MAX_BUCKET_DIM 96

// Static bucket boundaries (could also be __constant__ memory).
// Tercile cut points: 43, 86, 128 (sum = HIDDEN_DIM = 128).
__device__ const unsigned int BUCKET_OFFSETS[N_HORIZONS + 1] = {0, 43, 86, 128};

// ─────────────────────────────────────────────────────────────────────
// tau_sort_kernel: bitonic-merge sort of cfc.tau values.
//
// Launch: 1 block × 32 threads, shared_mem_bytes >= HIDDEN_DIM * 8 bytes
// (4 for tau key, 4 for original index value).
//
// Sorts ascending. Output: sorted_indices_d[p] = original channel index
// whose tau is the p-th smallest.
//
// Algorithm: bitonic merge sort over 128 elements, single block, 4 passes
// of 32 threads cooperating. Each thread handles 4 elements per pass.
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void tau_sort_kernel(
    const float* __restrict__ tau,    // [HIDDEN_DIM]
    unsigned int* __restrict__ sorted_indices  // [HIDDEN_DIM]
) {
    extern __shared__ float smem[];
    float* keys = smem;                              // [HIDDEN_DIM]
    unsigned int* vals = (unsigned int*)(smem + HIDDEN_DIM);  // [HIDDEN_DIM]

    int tid = threadIdx.x;
    // Each thread loads 4 elements (128 / 32 = 4)
    for (int i = 0; i < 4; ++i) {
        int idx = tid + i * 32;
        keys[idx] = tau[idx];
        vals[idx] = (unsigned int)idx;
    }
    __syncthreads();

    // Bitonic sort over 128 elements.
    for (int k = 2; k <= HIDDEN_DIM; k <<= 1) {
        for (int j = k >> 1; j > 0; j >>= 1) {
            for (int i = 0; i < 4; ++i) {
                int idx = tid + i * 32;
                int partner = idx ^ j;
                if (partner > idx) {
                    bool ascending = ((idx & k) == 0);
                    bool swap = ascending ? (keys[idx] > keys[partner]) : (keys[idx] < keys[partner]);
                    if (swap) {
                        float tk = keys[idx]; keys[idx] = keys[partner]; keys[partner] = tk;
                        unsigned int tv = vals[idx]; vals[idx] = vals[partner]; vals[partner] = tv;
                    }
                }
            }
            __syncthreads();
        }
    }

    // Write sorted_indices.
    for (int i = 0; i < 4; ++i) {
        int idx = tid + i * 32;
        sorted_indices[idx] = vals[idx];
    }
}

// ─────────────────────────────────────────────────────────────────────
// bucket_assign_kernel: from sorted_indices, write bucket_id per channel.
//
// Launch: 1 block × HIDDEN_DIM threads.
// Each thread handles one channel. Looks up its rank (position in sorted)
// and assigns bucket id via static tercile boundaries (post-2026-05-22
// N_HORIZONS=3 rebase: cuts at rank 43 and 86 for [43, 43, 42]).
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void bucket_assign_kernel(
    const unsigned int* __restrict__ sorted_indices,  // [HIDDEN_DIM]
    unsigned char* __restrict__ bucket_id_per_channel  // [HIDDEN_DIM]
) {
    int tid = threadIdx.x;
    if (tid >= HIDDEN_DIM) return;
    // tid is a sorted-rank position; sorted_indices[tid] is the original channel
    // at this rank. Assign bucket = tercile of rank.
    unsigned char bucket = 0;
    if (tid >= 86) bucket = 2;
    else if (tid >= 43) bucket = 1;
    bucket_id_per_channel[sorted_indices[tid]] = bucket;
}

// ─────────────────────────────────────────────────────────────────────
// bucket_iqr_kernel: per-bucket Q1 and Q3 of tau values.
//
// Launch: N_HORIZONS blocks × 32 threads. Each block computes Q1, Q3 for
// its bucket via warp-shuffle reduction.
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void bucket_iqr_kernel(
    const float* __restrict__ tau,                  // [HIDDEN_DIM]
    const unsigned int* __restrict__ sorted_indices, // [HIDDEN_DIM]
    float* __restrict__ q1_per_bucket,              // [N_HORIZONS]
    float* __restrict__ q3_per_bucket               // [N_HORIZONS]
) {
    int bucket = blockIdx.x;
    int bucket_start = BUCKET_OFFSETS[bucket];
    int bucket_end = BUCKET_OFFSETS[bucket + 1];
    int bucket_size = bucket_end - bucket_start;
    int q1_idx = bucket_start + bucket_size / 4;
    int q3_idx = bucket_start + (3 * bucket_size) / 4;
    if (threadIdx.x == 0) {
        // tau values are already in sorted-rank order via sorted_indices.
        q1_per_bucket[bucket] = tau[sorted_indices[q1_idx]];
        q3_per_bucket[bucket] = tau[sorted_indices[q3_idx]];
    }
}

// ─────────────────────────────────────────────────────────────────────
// channels_in_bucket_kernel: build the bucket → channel lookup table.
//
// ALPHA fix (2026-05-21): replaces `tau_reorder_kernel`. Per spec §2.3
// the original reorder approach grouped τ values by bucket position, but
// downstream kernels (`cfc_step_per_branch_fwd`, `tau_clamp_kernel`,
// `heads_w_skip_*_kernel`) all read using ORIGINAL channel indices —
// the reorder caused systematic layout mismatches.
//
// New approach: keep tau_all_d, W_in, W_rec, b, heads_w_skip in their
// ORIGINAL channel layout. At transition, materialise a lookup table
// `channels_in_bucket[k][i]` = original channel index of the i-th
// channel assigned to bucket k. Per-branch kernels iterate
// channels_in_bucket[branch][0..bucket_dim_k[branch]] to find the
// original channels in each bucket.
//
// Layout: `channels_in_bucket` is [N_HORIZONS × MAX_BUCKET_DIM] flat
// row-major. Sentinel value `0xFFFFFFFF` (UINT32_MAX) marks positions
// beyond `bucket_dim_k[k]` (the last bucket has 28 of 28 positions
// occupied; earlier buckets have 25 of 28).
//
// Launch: 1 block × HIDDEN_DIM threads, shared_mem_bytes = N_HORIZONS × 4
// for the per-bucket cursors. Single-thread sequential write (tid==0)
// keeps the cursor increments race-free without atomicAdd per
// feedback_no_atomicadd.
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void channels_in_bucket_kernel(
    const unsigned char* __restrict__ bucket_id_per_channel,  // [HIDDEN_DIM]
    const unsigned int* __restrict__ bucket_dim_k,            // [N_HORIZONS]
    unsigned int* __restrict__ channels_in_bucket             // [N_HORIZONS × MAX_BUCKET_DIM] flat
) {
    __shared__ unsigned int cursor[N_HORIZONS];
    int tid = threadIdx.x;

    // Init per-bucket cursors to 0.
    if (tid < N_HORIZONS) {
        cursor[tid] = 0;
    }
    __syncthreads();

    // Fill all output slots with the sentinel value first (parallel across
    // tid). For HIDDEN_DIM=128 ≥ N_HORIZONS*MAX_BUCKET_DIM=140 we use a
    // simple stride loop: each thread covers a few output positions.
    const int total_slots = N_HORIZONS * MAX_BUCKET_DIM;
    for (int idx = tid; idx < total_slots; idx += blockDim.x) {
        channels_in_bucket[idx] = 0xFFFFFFFFu;
    }
    __syncthreads();

    // Single-thread sequential write: thread 0 walks channels in original
    // order and appends each to the cursor position of its bucket. This
    // gives deterministic per-bucket ordering (channels listed in
    // ascending original-channel-index order). No atomicAdd, no race —
    // a single writer per feedback_no_atomicadd.
    if (tid == 0) {
        for (int c = 0; c < HIDDEN_DIM; ++c) {
            unsigned char k = bucket_id_per_channel[c];
            // Safety: bucket_id_per_channel must be in [0, N_HORIZONS).
            // bucket_assign_kernel guarantees this; the bound check below
            // is defensive only.
            if (k < (unsigned char)N_HORIZONS) {
                unsigned int pos = cursor[k];
                unsigned int bdim = bucket_dim_k[k];
                if (pos < bdim && pos < (unsigned int)MAX_BUCKET_DIM) {
                    channels_in_bucket[(unsigned int)k * MAX_BUCKET_DIM + pos] = (unsigned int)c;
                    cursor[k] = pos + 1;
                }
            }
        }
    }
}

// ─────────────────────────────────────────────────────────────────────
// zero_off_bucket_kernel — block-diagonal projection on
// [N_HORIZONS × HIDDEN_DIM] tensors.
//
// ALPHA fix (2026-05-21): the grad mask alone is insufficient to keep
// `heads_w_skip` block-diagonal because Adam's `m` and `v` momentum
// buffers carry pre-transition gradient signal across the boundary. The
// grad mask zeros NEW gradients at off-bucket positions; historical
// momentum continues to push them away from zero each Adam step.
//
// This kernel projects a generic `[N_HORIZONS × HIDDEN_DIM]` tensor
// onto the block-diagonal subspace by writing 0.0 at every (h, c) where
// `bucket_id_per_channel[c] != h`. It is launched on:
//   1. `heads_w_skip_d` at transition (one-shot, zeros off-bucket weights)
//   2. `heads_w_skip` Adam `m` and `v` buffers at transition (one-shot,
//      zeros off-bucket momentum so post-transition Adam steps cannot
//      revive off-bucket positions)
//   3. `heads_w_skip_d` after every Phase 2 Adam step (belt-and-suspenders;
//      guarantees off-bucket positions stay exactly zero even under
//      eps-driven Adam step from zero moments)
//
// Launch: grid = (N_HORIZONS, 1, 1), block = (HIDDEN_DIM, 1, 1). Same
// layout convention as `heads_w_skip_mask_init_kernel`.
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void zero_off_bucket_kernel(
    const unsigned char* __restrict__ bucket_id_per_channel,  // [HIDDEN_DIM]
    float* __restrict__ tensor                                 // [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;
    if (bucket_id_per_channel[c] != (unsigned char)h) {
        tensor[idx] = 0.0f;
    }
}

// ─────────────────────────────────────────────────────────────────────
// heads_compact_kernel: reorder heads_w_skip into compact ragged layout.
//
// Original heads_w_skip is [N_HORIZONS × HIDDEN_DIM]. Output is compact
// [HIDDEN_DIM] where each horizon h owns BUCKET_DIM[h] contiguous floats
// at position bucket_channel_offset[h].
//
// For each compact-position c in bucket k, read original
// heads_w_skip[k * HIDDEN_DIM + channel_that_landed_at_c].
//
// Launch: 1 block × HIDDEN_DIM threads.
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void heads_compact_kernel(
    const float* __restrict__ w_skip_orig,             // [N_HORIZONS × HIDDEN_DIM]
    const unsigned char* __restrict__ bucket_id_per_channel,  // [HIDDEN_DIM]
    const unsigned int* __restrict__ bucket_channel_offset,    // [N_HORIZONS + 1]
    float* __restrict__ w_skip_compact                 // [HIDDEN_DIM]
) {
    extern __shared__ unsigned int per_bucket_cursor[]; // [N_HORIZONS + 1]
    int tid = threadIdx.x;
    if (tid <= N_HORIZONS) {
        per_bucket_cursor[tid] = bucket_channel_offset[tid];
    }
    __syncthreads();

    for (int c = 0; c < HIDDEN_DIM; ++c) {
        if (tid == 0) {
            unsigned char k = bucket_id_per_channel[c];
            unsigned int pos = per_bucket_cursor[k];
            w_skip_compact[pos] = w_skip_orig[(unsigned int)k * HIDDEN_DIM + c];
            per_bucket_cursor[k] = pos + 1;
        }
        __syncthreads();
    }
}

// ─────────────────────────────────────────────────────────────────────
// tau_change_frobenius_kernel: ||tau_t - tau_{t-1}||_F² (scalar output).
//
// Single block × HIDDEN_DIM threads. Block-tree reduction (no atomicAdd
// per feedback_no_atomicadd).
//
// Per pearl_first_observation_bootstrap: caller's ControllerA bootstraps
// the EMA from this scalar on step 1; on subsequent steps the scalar
// feeds the Wiener-α update.
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void tau_change_frobenius_kernel(
    const float* __restrict__ tau_t,            // [HIDDEN_DIM]
    const float* __restrict__ tau_t_minus_1,    // [HIDDEN_DIM]
    float* __restrict__ tau_change_out          // [1]
) {
    __shared__ float sdata[HIDDEN_DIM];
    int tid = threadIdx.x;
    if (tid >= HIDDEN_DIM) return;
    float diff = tau_t[tid] - tau_t_minus_1[tid];
    sdata[tid] = diff * diff;
    __syncthreads();
    // Block-tree reduction (HIDDEN_DIM=128 is power-of-two; no padding needed).
    for (int s = HIDDEN_DIM / 2; s > 0; s >>= 1) {
        if (tid < s) sdata[tid] += sdata[tid + s];
        __syncthreads();
    }
    if (tid == 0) *tau_change_out = sdata[0];
}

// ─────────────────────────────────────────────────────────────────────
// slack_factor_apply_kernel: convert (Q1, Q3) → IQR_widened in place.
//
// Per spec §3.2: slack_factor_k = sqrt(Q3_k / Q1_k), ISV-derived from
// observed bucket IQR (no hardcoded 1.5× factor).
//
// Launch: 1 block × N_HORIZONS threads (5 threads, well within a warp).
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void slack_factor_apply_kernel(
    float* __restrict__ iqr_lo,   // [N_HORIZONS] in: Q1, out: Q1/slack
    float* __restrict__ iqr_hi    // [N_HORIZONS] in: Q3, out: Q3*slack
) {
    int k = threadIdx.x;
    if (k >= N_HORIZONS) return;
    float q1 = iqr_lo[k];
    float q3 = iqr_hi[k];
    // Numerical defense: Q1 must be > 0 (taus are log-uniform [0.01, 1000],
    // all positive by construction).
    float slack = sqrtf(q3 / fmaxf(q1, 1e-6f));
    iqr_lo[k] = q1 / slack;
    iqr_hi[k] = q3 * slack;
}

// ─────────────────────────────────────────────────────────────────────
// tau_clamp_kernel — Controller B post-Adam projection.
//
// Per spec §3.2 and pearl_adam_normalizes_loss_weights: hard projection
// AFTER each Adam step on the τ slice, bypassing Adam's m/sqrt(v)
// normalization that would no-op a loss-weight modulation.
//
// Per pearl_isv_for_adaptive_bounds: the (lo, hi) inputs are the IQR-
// widened bounds from `slack_factor_apply_kernel` (ISV-derived from each
// bucket's observed Q3/Q1 ratio at the Phase 1→2 transition).
//
// Launch: 1 block × HIDDEN_DIM threads (single warp on Ampere+; tiny).
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void tau_clamp_kernel(
    float* __restrict__ tau_all,                              // [HIDDEN_DIM]
    const unsigned char* __restrict__ bucket_id_per_channel,  // [HIDDEN_DIM]
    const float* __restrict__ iqr_lo,                         // [N_HORIZONS]
    const float* __restrict__ iqr_hi                          // [N_HORIZONS]
) {
    int c = threadIdx.x;
    if (c >= HIDDEN_DIM) return;
    unsigned char k = bucket_id_per_channel[c];
    float v = tau_all[c];
    float lo = iqr_lo[k];
    float hi = iqr_hi[k];
    tau_all[c] = fminf(fmaxf(v, lo), hi);
}

// ─────────────────────────────────────────────────────────────────────
// heads_lr_multiplier_scale_kernel — Controller C per-bucket LR multiplier.
//
// Per spec §3.3 and pearl_adam_normalizes_loss_weights: rescales the
// per-step parameter delta (post-Adam − pre-Adam) by a per-horizon LR
// multiplier. Bypasses Adam's m/sqrt(v) normalization because the
// rescaling is applied to the parameter delta itself, not the loss
// weight or the Adam moments.
//
// Layout: `param` is row-major [N_HORIZONS × per_horizon_stride]. Each
// horizon owns a contiguous block of `per_horizon_stride` floats. Element
// `i` belongs to horizon `i / per_horizon_stride`.
//
// Launch: enough threads to cover `total_elements` (one per element).
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void heads_lr_multiplier_scale_kernel(
    float* __restrict__ param,                       // [N_HORIZONS * per_horizon_stride]
    const float* __restrict__ param_pre,             // pre-Adam snapshot, same shape
    const float* __restrict__ lr_mult_per_horizon,   // [N_HORIZONS]
    int per_horizon_stride,
    int total_elements
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= total_elements) return;
    int horizon = i / per_horizon_stride;
    if (horizon >= N_HORIZONS) return;  // safety bound for stride/total mismatch
    float delta = param[i] - param_pre[i];
    float scaled_delta = delta * lr_mult_per_horizon[horizon];
    param[i] = param_pre[i] + scaled_delta;
}

// ─────────────────────────────────────────────────────────────────────
// h_mag_per_bucket_kernel — per-bucket mean(|h_state|) for Controller D.
//
// Per spec §3.4: Controller D's dead-bucket detector compares per-bucket
// `h_mag_k = mean(|h_state_branch_k|)` (across the batch × bucket_dim_k
// channels) against the first-observation floor with a 1/e decay
// threshold. This kernel computes the per-bucket mean-abs reduction on
// device; the host reads the [N_HORIZONS] result via a mapped-pinned
// shadow each Phase 2 step.
//
// ALPHA fix (2026-05-21): h_state stays in ORIGINAL channel layout (no
// reorder). The bucket → channel lookup goes through
// `channels_in_bucket[bucket][i]` populated at transition by
// `channels_in_bucket_kernel`.
//
// Launch: grid = (N_HORIZONS, 1, 1), block = (32, 1, 1). Each block
// reduces one bucket. Per `feedback_no_atomicadd`, reduction is
// block-tree on shared memory (no atomicAdd).
//
// Input `h_state` is `[B × HIDDEN_DIM]` (ORIGINAL channel layout). Each
// block reads its bucket's `bucket_dim_k[bucket]` channels (via
// `channels_in_bucket[bucket][0..bdim]`) per sample (across all B
// samples), accumulates |value|, and writes the per-bucket mean.
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void h_mag_per_bucket_kernel(
    const float* __restrict__ h_state,                          // [B × HIDDEN_DIM]
    const unsigned int* __restrict__ channels_in_bucket,        // [N_HORIZONS × MAX_BUCKET_DIM]
    const unsigned int* __restrict__ bucket_dim_k,              // [N_HORIZONS]
    int B,
    float* __restrict__ h_mag_per_bucket                        // [N_HORIZONS]
) {
    int bucket = blockIdx.x;
    if (bucket >= N_HORIZONS) return;
    int bdim = (int)bucket_dim_k[bucket];

    // sdata sized for a single-warp reduction (32 lanes).
    __shared__ float sdata[32];
    int tid = threadIdx.x;

    // Each thread sums |h| over its bucket-local index (tid), looking up
    // the ORIGINAL channel via channels_in_bucket. Threads with tid >= bdim
    // idle — uniform predicate, no warp divergence inside [0, 32).
    float local_sum = 0.0f;
    if (tid < bdim) {
        unsigned int c = channels_in_bucket[bucket * MAX_BUCKET_DIM + tid];
        // Defensive: sentinel slot OR out-of-range channel index should
        // not contribute. Predicate is uniform across the warp because all
        // active threads (tid < bdim) hold valid entries by construction
        // of channels_in_bucket_kernel.
        if (c < (unsigned int)HIDDEN_DIM) {
            for (int b = 0; b < B; ++b) {
                float v = h_state[b * HIDDEN_DIM + c];
                local_sum += fabsf(v);
            }
        }
    }
    sdata[tid] = (tid < 32) ? local_sum : 0.0f;
    __syncthreads();

    // Block-tree reduction over 32 lanes (single warp). No atomicAdd.
    for (int s = 16; s > 0; s >>= 1) {
        if (tid < s) sdata[tid] += sdata[tid + s];
        __syncthreads();
    }

    // Lane 0 writes the per-bucket mean. Total elements summed = bdim * B.
    if (tid == 0) {
        float denom = (float)(bdim * B);
        h_mag_per_bucket[bucket] = sdata[0] / fmaxf(denom, 1.0f);
    }
}

// ─────────────────────────────────────────────────────────────────────
// 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.
//
// Per spec §3.4: when a bucket's h_mag EMA falls below the first-
// observation × 1/e threshold for `dead_window_k` consecutive Phase 2
// steps, Recovery Attempt 1 doubles the bucket's slack envelope by
// halving its IQR lower bound and doubling the IQR upper bound. This
// releases τ values for that bucket from Controller B's hard projection
// clip, letting them drift to the wider range that may re-engage the
// channels.
//
// Per `feedback_isv_for_adaptive_bounds`: the factor 2 is the natural
// log-doubling step — one full log-spread expansion beyond the bucket's
// natural Q3/Q1 ratio (the IQR already encoded one log-step via the
// transition's `sqrt(Q3/Q1)` slack widening; doubling captures one more).
//
// Launch: grid = (1, 1, 1), block = (1, 1, 1). Single-thread, off the
// hot path (fires at most a handful of times per training run when a
// bucket actually goes dead).
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void bucket_iqr_double_widening_kernel(
    float* __restrict__ iqr_lo,   // [N_HORIZONS]
    float* __restrict__ iqr_hi,   // [N_HORIZONS]
    int bucket_idx
) {
    if (threadIdx.x == 0 && blockIdx.x == 0) {
        if (bucket_idx >= 0 && bucket_idx < N_HORIZONS) {
            iqr_lo[bucket_idx] *= 0.5f;
            iqr_hi[bucket_idx] *= 2.0f;
        }
    }
}