/*
 * aux_heads_kernel.cu — Plan 4 Task 6 Commit A; SP13 B1.1a (2026-05-05)
 * flipped the next-bar head from K=1 MSE regression to K=2 softmax CE
 * classification. The remaining brief comments describing K_out = 1 / MSE
 * are preserved BELOW where the rewrite history is relevant; the active
 * design is documented at the top of each kernel.
 *
 * Multi-task auxiliary heads (E.6): two small MLPs branching off the
 * trunk's `h_s2 [B, SH2]` post-GRN activation:
 *
 *   (1) Next-bar direction classification head (SP13 B1.1a):
 *         logits = Linear_2(ELU(Linear_1(h_s2)))   where K_out = 2
 *         softmax = stable_softmax(logits)
 *       Loss = cross-entropy against per-bar i32 direction label
 *              (0 = down, 1 = up, -1 = mask/skip). Producer wires in B1.1b.
 *
 *   (2) Five-class regime classification head:
 *         logits = Linear_2(ELU(Linear_1(h_s2)))   where K_out = 5
 *       Loss = cross-entropy against a discretised regime label
 *              (`floor(states[regime_score_idx] * 5)`, clamped to [0, 4]).
 *
 * Both heads share the `Linear(SH2 → 32) → ELU → Linear(32 → K)` shape
 * with `H = AUX_HIDDEN_DIM = 32`.
 *
 * SP13 B1.1a (2026-05-05): the next-bar head's K_out flipped 1 → 2 and
 * the loss formulation flipped MSE → softmax CE. The forward kernel now
 * emits BOTH the [B, K] logits tile (saved for backward via mean-shift
 * recompute is unnecessary — we save the softmax tile directly) AND a
 * [B, K] softmax tile. Three downstream consumers read the softmax tile:
 * the loss reduce, the backward kernel, and the dir-acc / isv-tanh
 * producers. Labels are i32 in [0, K) with `-1` reserved for "mask/skip".
 * Mean-over-batch CE divides by B_valid (count of non-masked rows), not
 * B, so an all-skip batch produces loss = 0 and zero gradients without
 * NaN. B1.1b lands the producer kernel that fills the labels.
 *
 * Pearls applied:
 *   - feedback_no_atomicadd.md — per-block shmem-tree reductions, no
 *     atomicAdd anywhere (param-grad accumulation is two-phase: per-block
 *     partial → final reduce, mirroring `grn_layernorm_backward_dgamma_dbeta`).
 *   - pearl_cold_path_no_exception_to_gpu_drives.md — every reduction
 *     stays GPU-side; no DtoH / host roundtrips.
 *   - feedback_no_feature_flags.md — no enable_/use_ booleans; the only
 *     axis of behavioural variation is K_out (1 vs 5), which is a
 *     structural property of the head, not a toggle.
 *
 * Layout convention: row-major `[B, …]` for all activation / gradient
 * tensors. `w1 [H, SH2]` row-major; `w2 [K, H]` row-major. Matches the
 * VSN per-group MLP convention from `vsn_feature_selection_kernel.cu`.
 *
 * ELU activation: f(x) = x if x > 0 else (exp(x) - 1).
 *   Backward: f'(x) = 1 if x > 0 else exp(x) = 1 + f(x).
 *
 * Cross-entropy: numerically-stable form — subtract max(logits) before exp.
 *   Loss[b]      = -log( exp(logits[b, label[b]] - max) / sum_k exp(logits[b, k] - max) )
 *   d_logits[b, k] = (softmax[b, k] - one_hot(label[b], k))
 *
 * Mean-over-batch loss reduction: dW / db both divide by B in the
 * accumulation reduce so the kernel-internal scale matches `1/B * sum_b ∂L/∂W`.
 */

#include <cuda_runtime.h>
#include <math_constants.h>

/* AUX_HIDDEN_DIM kept as a compile-time literal so kernel launch configs
 * don't need to thread `H` through. Both heads share the same hidden width. */
#ifndef AUX_HIDDEN_DIM
/* 2026-05-13: enlarged 32 → 128. Initial 256 attempt OOM'd L40S due
 * to per-sample dW1_partial [B, H, SH2] = 16384 × 256 × 256 × 4 = 4GB
 * per head. At H=128 it's 2GB per head = 4GB total, fits in 48GB
 * budget alongside next_states ~4GB. See gpu_aux_heads.rs for full
 * rationale. */
#define AUX_HIDDEN_DIM 128
#endif

/* Block dim used by reduction-shaped launches (per-tensor backward, label
 * builders). Kept identical to `vsn_feature_selection_kernel.cu`'s 256
 * threads/block convention. */
#ifndef AUX_BLOCK
#define AUX_BLOCK 256
#endif

/* ELU helpers — branchless on x sign via `x > 0.f`. */
__device__ __forceinline__ float aux_elu_fwd(float x) {
    return (x > 0.0f) ? x : (expf(x) - 1.0f);
}

/* Backward derivative from POST-activation y = ELU(x). For x > 0,
 * y == x and y > 0; for x <= 0, y = exp(x) - 1 ∈ (-1, 0] so 1 + y = exp(x).
 * Distinguishing on `y > 0` recovers the same partition. */
__device__ __forceinline__ float aux_elu_bwd_from_post(float y) {
    return (y > 0.0f) ? 1.0f : (1.0f + y);
}

/* =====================================================================
 * aux_next_bar_forward — Linear → ELU → Linear with K-way output (SP13
 * B1.1a: K_out = AUX_NEXT_BAR_K = 2, softmax classification).
 *
 * Per sample b (one block per sample, AUX_BLOCK threads):
 *   1. h[k] = ELU( b1[k] + sum_j w1[k, j] * h_s2[b, j] )   for k ∈ [0, H)
 *   2. logits[b, kc] = b2[kc] + sum_k w2[kc, k] * h[k]     for kc ∈ [0, K)
 *   3. softmax[b, kc] = stable_softmax(logits[b, :])[kc]    (max-shift)
 *
 * Hidden activations `h [B, H]` are SAVED for backward (consumed as the
 * ELU post-activation buffer; backward derives `f'(x_pre)` from `h_post`
 * via the standard ELU backward identity). The softmax tile `[B, K]` is
 * also saved — backward re-reads it (rather than recomputing from logits)
 * because the softmax tile has THREE downstream consumers in B1.1a:
 *
 *   - aux_next_bar_loss_reduce  (CE numerator: -log(softmax[label]))
 *   - aux_next_bar_backward     (d_logits = softmax - one_hot(label))
 *   - aux_dir_acc_reduce_kernel (argmax-vs-label compare)
 *   - aux_pred_to_isv_tanh_kernel (mean(softmax[1] - softmax[0]) → ISV[375])
 *
 * SP14 Layer C Phase C.5b (2026-05-08): the input parameter was renamed
 * `h_s2` → `h_s2_aux`. The aux heads now read the SEPARATE aux trunk's
 * output rather than Q's GRN trunk output. The Rust caller passes the
 * trainer's `h_s2_aux` buffer (populated by `aux_trunk_forward` per
 * `gpu_aux_trunk.rs`). Q's `save_h_s2` is no longer wired here. Symmetric
 * rename in `aux_next_bar_backward` (`dh_s2_out` → `dh_s2_aux_out`) so
 * the gradient flows back through the aux trunk to its own params, NOT
 * into Q's encoder trunk. The encoder boundary stop-gradient is enforced
 * structurally by `aux_trunk_backward`'s missing `dx_in` output param.
 *
 * Mirroring `aux_regime_forward`'s K-fanout structure — thread 0 computes
 * all K logits serially (K=2 / K=5; the matmul is small enough that a
 * tree reduce wastes more cycles in shfl bookkeeping than the serial
 * pass), then computes the stable softmax in-place on the same K-vector.
 *
 * Block: AUX_BLOCK threads. Grid: (B, 1, 1). Shared memory: H + K floats
 * (h cache for Linear_2 + temporary K-vector for stable softmax).
 * ===================================================================== */
extern "C" __global__ void aux_next_bar_forward(
    const float* __restrict__ h_s2_aux,    /* [B, SH2] row-major (SP14-C.5b: aux trunk output, NOT Q-trunk h_s2) */
    const float* __restrict__ w1,          /* [H, SH2] row-major */
    const float* __restrict__ b1,          /* [H] */
    const float* __restrict__ w2,          /* [K, H] row-major */
    const float* __restrict__ b2,          /* [K] */
    int B,
    int SH2,
    int K,
    float* __restrict__ hidden_out,        /* [B, H] OUT (saved for backward) */
    float* __restrict__ logits_out,        /* [B, K] OUT (saved for backward) */
    float* __restrict__ softmax_out        /* [B, K] OUT (saved for 3 consumers) */
) {
    const int b   = blockIdx.x;
    if (b >= B) return;
    const int tid = threadIdx.x;
    const int H   = AUX_HIDDEN_DIM;

    extern __shared__ float smem[];        /* [H + K] floats */
    float* sh_h     = smem;                /* [H] hidden cache */
    float* sh_logit = smem + H;            /* [K] logits cache for softmax */

    /* Step 1: per-hidden-unit Linear_1 + ELU. Each thread computes one or
     * more hidden lanes via stride loop. SH2 is small enough that each
     * thread doing a serial dot over SH2 is cheaper than block-cooperative
     * reduce per lane. */
    for (int k = tid; k < H; k += blockDim.x) {
        float acc = b1[k];
        const float* w_row = w1 + (size_t)k * SH2;
        const float* h_row = h_s2_aux + (size_t)b * SH2;
        for (int j = 0; j < SH2; ++j) {
            acc += w_row[j] * h_row[j];
        }
        const float h_post = aux_elu_fwd(acc);
        sh_h[k] = h_post;
        hidden_out[(size_t)b * H + k] = h_post;
    }
    __syncthreads();

    /* Step 2: K-way Linear_2. Thread 0 emits all K logits (K is small —
     * 2 in B1.1a; mirrors the regime kernel's K=5 fanout shape). Stores
     * each logit into the saved buffer AND a shmem scratch for Step 3's
     * softmax. */
    if (tid == 0) {
        for (int kc = 0; kc < K; ++kc) {
            float acc = b2[kc];
            const float* w2_row = w2 + (size_t)kc * H;
            for (int k = 0; k < H; ++k) {
                acc += w2_row[k] * sh_h[k];
            }
            logits_out[(size_t)b * K + kc] = acc;
            sh_logit[kc] = acc;
        }
    }
    __syncthreads();

    /* Step 3: stable softmax on the K-vector — subtract max before exp.
     * Thread 0 again (K=2; tree-reducing a K=2 vector is wasted bookkeeping). */
    if (tid == 0) {
        float lmax = sh_logit[0];
        for (int kc = 1; kc < K; ++kc) {
            if (sh_logit[kc] > lmax) lmax = sh_logit[kc];
        }
        float sum_e = 0.0f;
        for (int kc = 0; kc < K; ++kc) {
            sum_e += expf(sh_logit[kc] - lmax);
        }
        const float inv_sum = 1.0f / sum_e;
        for (int kc = 0; kc < K; ++kc) {
            softmax_out[(size_t)b * K + kc] = expf(sh_logit[kc] - lmax) * inv_sum;
        }
    }
}

/* =====================================================================
 * aux_regime_forward — Linear → ELU → Linear with K=5 output.
 *
 * Same architecture as next_bar but with K_out = 5.
 *
 * Per sample b: same Step 1 as next_bar (cached in shmem); Step 2 computes
 * 5 logits via thread 0 (5 × H = 160 multiplies — well under the cost of
 * a 5-way fanout block reduce).
 *
 * Block: AUX_BLOCK threads. Grid: (B, 1, 1). Shared memory: H floats.
 * ===================================================================== */
extern "C" __global__ void aux_regime_forward(
    const float* __restrict__ h_s2_aux,    /* [B, SH2] row-major (SP14-C.5b: aux trunk output, NOT Q-trunk h_s2) */
    const float* __restrict__ w1,          /* [H, SH2] row-major */
    const float* __restrict__ b1,          /* [H] */
    const float* __restrict__ w2,          /* [K, H] row-major (K = 5) */
    const float* __restrict__ b2,          /* [K] */
    int B,
    int SH2,
    int K,
    float* __restrict__ hidden_out,        /* [B, H] OUT (saved for backward) */
    float* __restrict__ logits_out         /* [B, K] OUT */
) {
    const int b   = blockIdx.x;
    if (b >= B) return;
    const int tid = threadIdx.x;
    const int H   = AUX_HIDDEN_DIM;

    extern __shared__ float smem[]; /* h cache [H] */

    for (int k = tid; k < H; k += blockDim.x) {
        float acc = b1[k];
        const float* w_row = w1 + (size_t)k * SH2;
        const float* h_row = h_s2_aux + (size_t)b * SH2;
        for (int j = 0; j < SH2; ++j) {
            acc += w_row[j] * h_row[j];
        }
        const float h_post = aux_elu_fwd(acc);
        smem[k] = h_post;
        hidden_out[(size_t)b * H + k] = h_post;
    }
    __syncthreads();

    /* Step 2: K-way Linear_2. Thread 0 emits all K logits (K is at most 5
     * in this commit's design). */
    if (tid == 0) {
        for (int kc = 0; kc < K; ++kc) {
            float acc = b2[kc];
            const float* w2_row = w2 + (size_t)kc * H;
            for (int k = 0; k < H; ++k) {
                acc += w2_row[k] * smem[k];
            }
            logits_out[(size_t)b * K + kc] = acc;
        }
    }
}

/* =====================================================================
 * aux_next_bar_loss_reduce — single-block shmem-reduce mean cross-entropy
 * over the K-class softmax tile (SP13 B1.1a; flipped from MSE).
 *
 * For each sample b:
 *   if labels[b] == -1: contribute 0 to numerator AND 0 to valid count
 *   else              : numerator += -log(softmax[b, labels[b]])
 *                       valid_count += 1
 *
 * Reduces both the numerator and the valid count in lockstep via two
 * parallel shmem-tree strips, then writes
 *   loss_out[0]        = numerator / max(valid_count, 1.0f)
 *   valid_count_out[0] = valid_count   (saved for backward)
 *
 * Mean-over-valid-bars semantics: an all-skip batch produces loss = 0
 * (no NaN from `0 / max(0, 1)` since the divisor is masked to ≥1) and
 * the saved valid_count = 0 makes the backward kernel emit zero
 * gradients across the board.
 *
 * Single-block (AUX_BLOCK threads) shmem-tree reduction; no atomicAdd
 * per `feedback_no_atomicadd.md`. Two parallel partial-reduction strips
 * (loss numerator + valid count); shmem layout matches the K-class
 * `aux_regime_loss_reduce` shape.
 *
 * The numerically-stable softmax was already materialised in the forward
 * kernel — this reduce just reads `softmax[b, labels[b]]` directly,
 * avoiding a second max-shift+exp pass per row.
 *
 * Block: AUX_BLOCK threads. Grid: (1, 1, 1).
 * Shared memory: 2 * AUX_BLOCK floats (loss + valid-count partials).
 * ===================================================================== */
extern "C" __global__ void aux_next_bar_loss_reduce(
    const float* __restrict__ softmax,         /* [B, K] row-major */
    const int* __restrict__ labels,            /* [B] i32 in {-1, 0, 1, …, K-1} */
    int B,
    int K,
    float* __restrict__ loss_out,              /* [1] mean CE over valid rows */
    float* __restrict__ valid_count_out        /* [1] B_valid (read by backward) */
) {
    if (blockIdx.x != 0) return;
    extern __shared__ float smem[];
    float* sh_loss  = smem;
    float* sh_valid = smem + blockDim.x;
    const int tid   = threadIdx.x;
    const int block = blockDim.x;

    /* Per-thread strided accumulation: each thread folds its slice of
     * `[0, B)` into private (loss_numer, valid_count) accumulators.
     * Mask-encoded rows (label == -1) contribute 0 to BOTH accumulators —
     * no branch in the loss path; the divisor uses fmaxf(valid, 1) to
     * prevent NaN when the entire batch is masked. */
    float local_loss  = 0.0f;
    float local_valid = 0.0f;
    for (int i = tid; i < B; i += block) {
        const int tgt = labels[i];
        if (tgt == -1) continue;
        /* Defensive bounds: out-of-range labels degrade to skip rather
         * than read past the end of the [B, K] softmax row. The producer
         * (B1.1b) is responsible for emitting only -1 / [0, K) — this
         * branch is dead under correct producer behavior, kept for kernel
         * robustness against producer bugs. */
        if (tgt < 0 || tgt >= K) continue;
        const float p_tgt = softmax[(size_t)i * K + tgt];
        /* Numerical floor: the forward softmax can produce p ≈ 0 for
         * extreme negative logits; -log(0) = +inf would propagate to the
         * EMA. Floor at FLT_MIN-ish (1e-30); equivalent to clipping the
         * logit at log(1e-30) ≈ -69 which is well outside the trained
         * regime. */
        const float p_clipped = fmaxf(p_tgt, 1e-30f);
        local_loss  += -logf(p_clipped);
        local_valid += 1.0f;
    }
    sh_loss[tid]  = local_loss;
    sh_valid[tid] = local_valid;
    __syncthreads();

    /* Standard log2(BLOCK) tree reduction over both arrays in lockstep. */
    for (int s = block / 2; s > 0; s >>= 1) {
        if (tid < s) {
            sh_loss[tid]  += sh_loss[tid + s];
            sh_valid[tid] += sh_valid[tid + s];
        }
        __syncthreads();
    }

    if (tid == 0) {
        const float valid = sh_valid[0];
        loss_out[0]        = sh_loss[0] / fmaxf(valid, 1.0f);
        valid_count_out[0] = valid;
    }
}

/* =====================================================================
 * aux_regime_loss_reduce — single-block shmem-reduce mean cross-entropy.
 *
 * For each sample b, computes per-sample CE = -log(softmax(logits)[label])
 * using the numerically-stable max-shift form. Reduces both the loss and
 * the count of samples whose argmax(logits) == label, writing scalar
 * means into loss_out[0] and correct_out[0].
 *
 * Block: AUX_BLOCK threads. Grid: (1, 1, 1).
 * Shared memory: 2 * AUX_BLOCK floats (loss + correct partial reductions).
 * ===================================================================== */
extern "C" __global__ void aux_regime_loss_reduce(
    const float* __restrict__ logits,      /* [B, K] row-major */
    const unsigned char* __restrict__ label,  /* [B] uint8 in [0, K) */
    int B,
    int K,
    float* __restrict__ loss_out,          /* [1] */
    float* __restrict__ correct_out        /* [1] mean accuracy in [0, 1] */
) {
    if (blockIdx.x != 0) return;
    extern __shared__ float smem[];
    float* sh_loss    = smem;
    float* sh_correct = smem + blockDim.x;
    const int tid   = threadIdx.x;
    const int block = blockDim.x;

    float local_loss    = 0.0f;
    float local_correct = 0.0f;

    for (int i = tid; i < B; i += block) {
        const float* row = logits + (size_t)i * K;

        /* Numerically-stable softmax: subtract max before exp. */
        float lmax = row[0];
        int   amax = 0;
        for (int k = 1; k < K; ++k) {
            if (row[k] > lmax) { lmax = row[k]; amax = k; }
        }
        float sum_e = 0.0f;
        for (int k = 0; k < K; ++k) {
            sum_e += expf(row[k] - lmax);
        }
        const int   tgt   = (int)label[i];
        const float ce    = (lmax - row[tgt]) + logf(sum_e);
        local_loss   += ce;
        local_correct += (amax == tgt) ? 1.0f : 0.0f;
    }
    sh_loss[tid]    = local_loss;
    sh_correct[tid] = local_correct;
    __syncthreads();

    for (int s = block / 2; s > 0; s >>= 1) {
        if (tid < s) {
            sh_loss[tid]    += sh_loss[tid + s];
            sh_correct[tid] += sh_correct[tid + s];
        }
        __syncthreads();
    }

    if (tid == 0) {
        loss_out[0]    = sh_loss[0]    / (float)B;
        correct_out[0] = sh_correct[0] / (float)B;
    }
}

/* =====================================================================
 * aux_regime_label_from_states — discretise states[:, regime_score_idx]
 * into uint8 labels in [0, K).
 *
 * Reads `states[b, regime_score_idx]` (a float in [0, 1] — the Plan 4
 * spec's regime stability score / regime cluster score), maps to an
 * integer bin via floor(s * K) clamped to [0, K-1].
 *
 * Block: AUX_BLOCK threads. Grid: ceil(B / AUX_BLOCK) blocks. One thread
 * per sample. No atomicAdd (each thread writes a unique label slot).
 * ===================================================================== */
extern "C" __global__ void aux_regime_label_from_states(
    const float* __restrict__ states,      /* [B, F] row-major */
    int B,
    int F,
    int regime_score_idx,
    int K,
    unsigned char* __restrict__ label_out  /* [B] uint8 */
) {
    const int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= B) return;

    const float s = states[(size_t)b * F + regime_score_idx];
    /* floor(s * K) with K in {1..255}; clamp to [0, K-1]. The clamp covers
     * both the s == 1.0 boundary (floor(K) == K out-of-range) and any
     * out-of-spec negative values. */
    int bin = (int)floorf(s * (float)K);
    if (bin < 0)      bin = 0;
    if (bin > K - 1)  bin = K - 1;
    label_out[b] = (unsigned char)bin;
}

/* =====================================================================
 * aux_next_bar_backward — full backward for the next-bar regression head.
 *
 * Forward was:
 *   h_pre[b, k] = b1[k] + sum_j w1[k, j] * h_s2[b, j]
 *   h_post[b, k] = ELU(h_pre[b, k])
 *   pred[b]      = b2[0] + sum_k w2[0, k] * h_post[b, k]
 *
 * Loss = (1/B) * sum_b (pred[b] - label[b])^2.
 *
 * Mean-over-batch backward:
 *   d_pred[b] = (2 / B) * (pred[b] - label[b])
 *   d_b2[0]  = sum_b d_pred[b]
 *   d_w2[0, k] = sum_b d_pred[b] * h_post[b, k]
 *   d_h_post[b, k] = d_pred[b] * w2[0, k]
 *   d_h_pre[b, k]  = d_h_post[b, k] * elu_bwd(h_post[b, k])
 *   d_b1[k]  = sum_b d_h_pre[b, k]
 *   d_w1[k, j] = sum_b d_h_pre[b, k] * h_s2[b, j]
 *   d_h_s2[b, j] = sum_k d_h_pre[b, k] * w1[k, j]
 *
 * Implementation: one block per sample (B blocks), AUX_BLOCK threads.
 * Each block computes its sample's `d_h_pre` into shared memory, then:
 *   - writes its sample's contribution to `d_h_s2` row b (no contention),
 *   - writes per-block partials `dW1_partial[b, k, j]` and
 *     `dW2_partial[b, 0, k]` and per-block bias partials.
 *
 * To stay atomicAdd-free, we DO NOT do block-cooperative param-grad
 * accumulation. Instead the caller runs a follow-up REDUCE pass over the
 * `[B, …]` partials. This kernel emits the per-sample partials as
 * contiguous tensors (size O(B * H * SH2 + B * H + B * H + B * 1)) and a
 * separate `aux_param_grad_reduce_*` kernel finalises them.
 *
 * For Commit A (no callers), the partial-emission contract is the
 * artefact that Commit B's training-loop wire-up consumes; the reduce
 * step lands together with the wire-up since the partials → final-reduce
 * pattern needs both endpoints to be exercised at once.
 *
 * The dh_s2 output is `[B, SH2]`: the trunk's h_s2 backward path will SUM
 * this into the existing main-path d_h_s2 (no atomicAdd, just a SAXPY in
 * the caller).
 *
 * Block: AUX_BLOCK threads. Grid: (B, 1, 1).
 * Shared memory: 2 * H floats (d_h_pre cache + ELU-derivative cache).
 * ===================================================================== */
extern "C" __global__ void aux_next_bar_backward(
    /* Forward inputs (re-read for dW). */
    const float* __restrict__ h_s2_aux,    /* [B, SH2] (SP14-C.5b: aux trunk output) */
    const float* __restrict__ w1,          /* [H, SH2] */
    const float* __restrict__ w2,          /* [K, H] row-major */
    const float* __restrict__ hidden_post, /* [B, H] saved h_post */
    /* Loss inputs (SP13 B1.1a softmax CE). */
    const float* __restrict__ softmax,     /* [B, K] saved softmax tile */
    const int*   __restrict__ labels,      /* [B] i32 in {-1, 0..K-1} */
    const float* __restrict__ valid_count, /* [1] B_valid from loss kernel */
    int B,
    int SH2,
    int K,
    /* Partial outputs (per-sample). Caller reduces along batch dim. */
    float* __restrict__ dW1_partial,       /* [B, H, SH2] flat */
    float* __restrict__ db1_partial,       /* [B, H] */
    float* __restrict__ dW2_partial,       /* [B, K, H] flat */
    float* __restrict__ db2_partial,       /* [B, K] */
    /* Aux trunk-output gradient: SP14-C.5b accumulator
     * `dh_s2_aux_accum` (NOT Q's `bw_d_h_s2`). Caller SAXPYs both heads'
     * dh_s2_aux into this buffer, then aux_trunk_backward reads it as the
     * upstream gradient feeding `dh_s2_aux_in_ptr`. Encoder stop-grad is
     * enforced structurally by aux_trunk_backward (no dx_in output). */
    float* __restrict__ dh_s2_aux_out      /* [B, SH2] */
) {
    const int b   = blockIdx.x;
    if (b >= B) return;
    const int tid = threadIdx.x;
    const int H   = AUX_HIDDEN_DIM;

    extern __shared__ float smem[];
    float* sh_dh_pre  = smem;           /* [H] */
    float* sh_h_post  = smem + H;       /* [H] */
    float* sh_dlogits = smem + 2 * H;   /* [K] */

    /* SP13 B1.1a (2026-05-05): softmax CE backward. Mirrors
     * `aux_regime_backward`'s K-fanout structure but reads the
     * pre-computed softmax tile from forward (rather than re-softmaxing
     * the logits) since the tile is already saved for three other
     * consumers. Mean-over-valid-rows derivative:
     *   d_logits[b, kc] = (softmax[b, kc] - one_hot(labels[b], kc)) / B_valid
     * for valid rows; masked rows (labels[b] == -1) zero out d_logits
     * across the K-vector so the partial-grad write becomes a no-op.
     * B_valid is read from the [1] scalar that the loss kernel wrote in
     * the same step; floor at 1.0f to mirror the loss-kernel's
     * fmaxf(valid, 1) divisor (an all-skip batch produces zero gradients
     * across the board — d_logits = 0 for every row, no NaN). */
    const int   tgt   = labels[b];
    const float B_v   = fmaxf(valid_count[0], 1.0f);
    const float inv_B = 1.0f / B_v;

    /* Step 0: compute d_logits[b, :] from softmax + label.
     * Single thread — K is small (2 in B1.1a; mirrors regime kernel
     * pattern at K=5). Masked rows zero the K-vector. */
    if (tid == 0) {
        if (tgt == -1 || tgt < 0 || tgt >= K) {
            for (int kc = 0; kc < K; ++kc) {
                sh_dlogits[kc] = 0.0f;
            }
        } else {
            const float* p_row = softmax + (size_t)b * K;
            for (int kc = 0; kc < K; ++kc) {
                const float onehot = (kc == tgt) ? 1.0f : 0.0f;
                sh_dlogits[kc] = inv_B * (p_row[kc] - onehot);
            }
        }
    }
    __syncthreads();

    /* Step 1: cache h_post and compute d_h_pre per hidden unit.
     *   d_h_post[b, k] = sum_kc d_logits[b, kc] * w2[kc, k]
     *   d_h_pre[b, k]  = d_h_post[b, k] * elu_bwd(h_post[b, k]) */
    for (int k = tid; k < H; k += blockDim.x) {
        const float h_post = hidden_post[(size_t)b * H + k];
        sh_h_post[k] = h_post;

        float d_h_post = 0.0f;
        for (int kc = 0; kc < K; ++kc) {
            d_h_post += sh_dlogits[kc] * w2[(size_t)kc * H + k];
        }
        sh_dh_pre[k] = d_h_post * aux_elu_bwd_from_post(h_post);
    }
    __syncthreads();

    /* Step 2: per-sample param-grad partials.
     *   db2_partial[b, kc]    = sh_dlogits[kc]
     *   dW2_partial[b, kc, k] = sh_dlogits[kc] * h_post[b, k]
     *   db1_partial[b, k]     = sh_dh_pre[k]
     *   dW1_partial[b, k, j]  = sh_dh_pre[k] * h_s2[b, j] */
    if (tid == 0) {
        for (int kc = 0; kc < K; ++kc) {
            db2_partial[(size_t)b * K + kc] = sh_dlogits[kc];
        }
    }

    /* dW2_partial[b, kc, k] = sh_dlogits[kc] * sh_h_post[k]. */
    {
        const int total = K * H;
        for (int idx = tid; idx < total; idx += blockDim.x) {
            const int kc = idx / H;
            const int k  = idx - kc * H;
            dW2_partial[(size_t)b * K * H + idx] = sh_dlogits[kc] * sh_h_post[k];
        }
    }

    for (int k = tid; k < H; k += blockDim.x) {
        db1_partial[(size_t)b * H + k] = sh_dh_pre[k];
    }

    /* dW1 partial — per-(k, j) write. Stride loop over (H * SH2). */
    {
        const float* h_row = h_s2_aux + (size_t)b * SH2;
        float* dW1_b = dW1_partial + (size_t)b * H * SH2;
        const int total = H * SH2;
        for (int idx = tid; idx < total; idx += blockDim.x) {
            const int k = idx / SH2;
            const int j = idx - k * SH2;
            dW1_b[idx] = sh_dh_pre[k] * h_row[j];
        }
    }

    /* Step 3: dh_s2_aux[b, j] = sum_k sh_dh_pre[k] * w1[k, j].
     *
     * SP14 Layer C Phase C.5b (2026-05-08): the prior commit `872bd7392`
     * zero-filled this buffer as a stop-grad band-aid because the aux
     * head's gradient was leaking into Q's shared `h_s2`. Phase C.5b
     * REVERTS that band-aid: aux heads now read the SEPARATE aux trunk's
     * output `h_s2_aux`, and `dh_s2_aux_out` accumulates into the aux
     * trunk's own input gradient (`dh_s2_aux_accum`), NOT into Q's
     * `bw_d_h_s2`. The aux trunk's backward kernel terminates at the
     * encoder boundary (no `dx_in` output param), so Q's encoder is
     * structurally protected — no leak possible. With that structural
     * guarantee the genuine SAXPY-back-to-input gradient is restored.
     *
     * Per-sample, per-feature serial reduction over H lanes — H is small
     * (32) and SH2 is moderate; stride loop over j. Masked rows propagate
     * zero d_h_pre through this sum, so dh_s2_aux row is zero for skipped
     * samples (the trunk SAXPY adds zero — no spurious gradient). */
    {
        float* dh_row = dh_s2_aux_out + (size_t)b * SH2;
        for (int j = tid; j < SH2; j += blockDim.x) {
            float acc = 0.0f;
            for (int k = 0; k < H; ++k) {
                acc += sh_dh_pre[k] * w1[(size_t)k * SH2 + j];
            }
            dh_row[j] = acc;
        }
    }
}

/* =====================================================================
 * aux_regime_backward — full backward for the 5-class regime CE head.
 *
 * Forward was:
 *   logits[b, kc] = b2[kc] + sum_k w2[kc, k] * h_post[b, k]
 *   p[b, kc]      = softmax(logits[b, :])[kc]
 *   CE[b]         = -log(p[b, label[b]])
 *
 * Mean-over-batch loss `L = (1/B) * sum_b CE[b]` ⇒
 *   d_logits[b, kc] = (1/B) * (p[b, kc] - one_hot(label[b], kc))
 *
 * The remaining backward (dW2, db2, dW1, db1, dh_s2) mirrors the structure
 * of `aux_next_bar_backward` but with a K-vector `d_logits` instead of a
 * scalar `d_pred`.
 *
 * Block: AUX_BLOCK threads. Grid: (B, 1, 1).
 * Shared memory: 2 * H + K floats (d_h_pre cache, h_post cache, d_logits cache).
 * ===================================================================== */
extern "C" __global__ void aux_regime_backward(
    /* Forward inputs. */
    const float* __restrict__ h_s2_aux,    /* [B, SH2] (SP14-C.5b: aux trunk output) */
    const float* __restrict__ w1,          /* [H, SH2] */
    const float* __restrict__ w2,          /* [K, H] row-major */
    const float* __restrict__ hidden_post, /* [B, H] saved h_post */
    const float* __restrict__ logits,      /* [B, K] saved (re-softmaxed here) */
    /* Loss inputs. */
    const unsigned char* __restrict__ label, /* [B] uint8 in [0, K) */
    int B,
    int SH2,
    int K,
    /* Partial outputs (per-sample). */
    float* __restrict__ dW1_partial,       /* [B, H, SH2] */
    float* __restrict__ db1_partial,       /* [B, H] */
    float* __restrict__ dW2_partial,       /* [B, K, H] */
    float* __restrict__ db2_partial,       /* [B, K] */
    /* Aux trunk-output gradient (SAXPYed into `dh_s2_aux_accum`; encoder
     * stop-grad enforced structurally by `aux_trunk_backward`). */
    float* __restrict__ dh_s2_aux_out      /* [B, SH2] */
) {
    const int b   = blockIdx.x;
    if (b >= B) return;
    const int tid = threadIdx.x;
    const int H   = AUX_HIDDEN_DIM;

    extern __shared__ float smem[];
    float* sh_dh_pre  = smem;             /* [H] */
    float* sh_h_post  = smem + H;         /* [H] */
    float* sh_dlogits = smem + 2 * H;     /* [K] */

    /* Step 0: compute d_logits[b, :] from the saved logits + label. Single
     * thread (K is small — at most 5). Mean-over-batch scaling: divide by B. */
    if (tid == 0) {
        const float* row = logits + (size_t)b * K;
        float lmax = row[0];
        for (int k = 1; k < K; ++k) if (row[k] > lmax) lmax = row[k];
        float sum_e = 0.0f;
        for (int k = 0; k < K; ++k) sum_e += expf(row[k] - lmax);
        const float inv_sum = 1.0f / sum_e;
        const int   tgt     = (int)label[b];
        const float inv_B   = 1.0f / (float)B;
        for (int k = 0; k < K; ++k) {
            const float p = expf(row[k] - lmax) * inv_sum;
            const float onehot = (k == tgt) ? 1.0f : 0.0f;
            sh_dlogits[k] = inv_B * (p - onehot);
        }
    }
    __syncthreads();

    /* Step 1: cache h_post and compute d_h_pre.
     *   d_h_post[b, k] = sum_kc d_logits[b, kc] * w2[kc, k]
     *   d_h_pre[b, k]  = d_h_post[b, k] * elu_bwd(h_post[b, k])
     */
    for (int k = tid; k < H; k += blockDim.x) {
        const float h_post = hidden_post[(size_t)b * H + k];
        sh_h_post[k] = h_post;

        float d_h_post = 0.0f;
        for (int kc = 0; kc < K; ++kc) {
            d_h_post += sh_dlogits[kc] * w2[(size_t)kc * H + k];
        }
        sh_dh_pre[k] = d_h_post * aux_elu_bwd_from_post(h_post);
    }
    __syncthreads();

    /* Step 2: per-sample param-grad partials. */
    if (tid == 0) {
        for (int kc = 0; kc < K; ++kc) {
            db2_partial[(size_t)b * K + kc] = sh_dlogits[kc];
        }
    }

    /* dW2_partial[b, kc, k] = sh_dlogits[kc] * sh_h_post[k]. */
    {
        const int total = K * H;
        for (int idx = tid; idx < total; idx += blockDim.x) {
            const int kc = idx / H;
            const int k  = idx - kc * H;
            dW2_partial[(size_t)b * K * H + idx] = sh_dlogits[kc] * sh_h_post[k];
        }
    }

    for (int k = tid; k < H; k += blockDim.x) {
        db1_partial[(size_t)b * H + k] = sh_dh_pre[k];
    }

    /* dW1_partial[b, k, j] = sh_dh_pre[k] * h_s2_aux[b, j]. */
    {
        const float* h_row = h_s2_aux + (size_t)b * SH2;
        float* dW1_b = dW1_partial + (size_t)b * H * SH2;
        const int total = H * SH2;
        for (int idx = tid; idx < total; idx += blockDim.x) {
            const int k = idx / SH2;
            const int j = idx - k * SH2;
            dW1_b[idx] = sh_dh_pre[k] * h_row[j];
        }
    }

    /* Step 3: dh_s2_aux[b, j] = sum_k sh_dh_pre[k] * w1[k, j].
     *
     * SP14 Layer C Phase C.5b (2026-05-08): the prior commit `411a30473`
     * zero-filled this buffer as a stop-grad band-aid because the regime
     * head's gradient was leaking into Q's shared `h_s2`. Phase C.5b
     * REVERTS that band-aid: regime now reads `h_s2_aux` (aux trunk
     * output), and `dh_s2_aux_out` SAXPYs into the aux trunk's own input
     * gradient `dh_s2_aux_accum`. The aux trunk's backward kernel has
     * NO `dx_in` output param so Q's encoder is structurally protected.
     * The genuine SAXPY-back-to-input gradient is restored. */
    {
        float* dh_row = dh_s2_aux_out + (size_t)b * SH2;
        for (int j = tid; j < SH2; j += blockDim.x) {
            float acc = 0.0f;
            for (int k = 0; k < H; ++k) {
                acc += sh_dh_pre[k] * w1[(size_t)k * SH2 + j];
            }
            dh_row[j] = acc;
        }
    }
}

/* =====================================================================
 * aux_param_grad_reduce — final-reduce kernel for the per-sample partials
 * emitted by both backward kernels.
 *
 * Aggregates `partials [B, P]` along the batch dimension into
 * `final [P]`. Single-block 256-thread (or grid-stride per-element)
 * shmem-tree reduction; no atomicAdd. Each block handles one column of
 * the partials (one element of the final output).
 *
 * Grid: (P, 1, 1) blocks. Block: AUX_BLOCK threads. Shared mem: AUX_BLOCK floats.
 * ===================================================================== */
extern "C" __global__ void aux_param_grad_reduce(
    const float* __restrict__ partials,    /* [B, P] row-major */
    int B,
    int P,
    float* __restrict__ final_out          /* [P] */
) {
    const int p = blockIdx.x;
    if (p >= P) return;
    const int tid   = threadIdx.x;
    const int block = blockDim.x;

    extern __shared__ float smem[];

    float local = 0.0f;
    for (int i = tid; i < B; i += block) {
        local += partials[(size_t)i * P + p];
    }
    smem[tid] = local;
    __syncthreads();

    for (int s = block / 2; s > 0; s >>= 1) {
        if (tid < s) smem[tid] += smem[tid + s];
        __syncthreads();
    }

    if (tid == 0) {
        final_out[p] = smem[0];
    }
}
