foxhunt/docs/superpowers/plans/2026-04-15-supervised-architecture-transfer.md

Supervised Architecture Transfer Implementation Plan

For agentic workers: REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (- [ ]) syntax for tracking.

Goal: Transfer 9 architectural concepts (7 from supervised suite + 2 stability fixes from live H100 observations) into the DQN training pipeline.

Architecture: 9 components. Tasks 10-11 (Q-mean centering + atom entropy recovery) should be executed FIRST — they fix active instabilities observed in live H100 run train-vpb4w (Q-mean drift 0→+0.47, atom utilization 100%→20%). Then the 7 supervised transfers deploy incrementally.

Tech Stack: Rust 1.85+, CUDA 12.4 (NVRTC precompiled cubins via build.rs), cuBLAS cublasLt TF32, cudarc driver bindings. All kernels in experience_kernels.cu, all Rust integration in the crates/ml/src/cuda_pipeline/ and crates/ml/src/trainers/dqn/ modules.

---

Task 1: CUDA Kernels -- All New Kernels

Files:

• crates/ml/src/cuda_pipeline/experience_kernels.cu

All 7 new CUDA kernels are added to experience_kernels.cu (the single source file that compiles to experience_kernels.cubin, which is loaded as EXPECTED_Q_CUBIN in gpu_dqn_trainer.rs). No build.rs changes needed -- the existing compilation entry at line 73 already covers this file.

• Step 1: Read crates/ml/src/cuda_pipeline/experience_kernels.cu from the end of the file (after the glu_backward kernel at ~line 2748) to find the insertion point.

• Step 2: Add the quantile_q_select kernel after the existing glu_backward kernel. This kernel extracts 10th/50th/90th percentile Q-values from C51 atom distributions and blends them using IQN readiness:

/* ================================================================== */
/* Kernel: quantile_q_select                                           */
/* ================================================================== */

/**
 * Uncertainty-driven action selection from C51 atom distributions.
 *
 * Extracts Q_10th, Q_50th, Q_90th from the CDF of each action's atom
 * distribution, then blends: Q_select = (1-alpha)*Q_90th + alpha*Q_10th
 * where alpha = iqn_readiness (0=exploring/optimistic, 1=exploiting/pessimistic).
 *
 * Replaces Boltzmann-on-E[Q] action selection in experience collection.
 * E[Q] is still computed separately for Q-stats and loss computation.
 *
 * Grid: ceil(N/256), Block: 256. One thread per episode.
 *
 * @param v_logits          [N, NA]                    value head logits
 * @param b_logits          [N, (B0+B1+B2+B3)*NA]      branch advantage logits (branch-major)
 * @param q_select          [N, B0+B1+B2+B3]           output: quantile-blended Q per action
 * @param N                                             number of samples
 * @param num_atoms                                     C51 atom count (NA)
 * @param b0_size                                       direction branch size
 * @param b1_size                                       magnitude branch size
 * @param b2_size                                       order branch size
 * @param b3_size                                       urgency branch size
 * @param per_sample_support [N*3]                      per-sample [v_min, v_max, delta_z]
 * @param iqn_readiness                                 scalar alpha for quantile blend
 */
extern "C" __global__ void quantile_q_select(
    const float* __restrict__ v_logits,
    const float* __restrict__ b_logits,
    float*                q_select,
    int N,
    int num_atoms,
    int b0_size,
    int b1_size,
    int b2_size,
    int b3_size,
    const float* __restrict__ per_sample_support,
    float iqn_readiness)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= N) return;

    float v_min = per_sample_support[i * 3];
    float dz = per_sample_support[i * 3 + 2];

    int total_actions = b0_size + b1_size + b2_size + b3_size;
    const float* v_row = v_logits + (long long)i * num_atoms;

    int branch_sizes[4] = { b0_size, b1_size, b2_size, b3_size };
    int branch_logit_offset = 0;
    int action_idx = 0;

    float alpha = fminf(fmaxf(iqn_readiness, 0.0f), 1.0f);

    for (int d = 0; d < 4; d++) {
        int n_d = branch_sizes[d];
        const float* branch_base = b_logits + (long long)branch_logit_offset
                                 + (long long)i * n_d * num_atoms;

        for (int a = 0; a < n_d; a++) {
            const float* adv_a = branch_base + (long long)a * num_atoms;

            /* Softmax over atoms */
            float max_logit = -1e30f;
            for (int z = 0; z < num_atoms; z++) {
                float logit = v_row[z] + adv_a[z];
                if (logit > max_logit) max_logit = logit;
            }
            float sum_exp = 0.0f;
            for (int z = 0; z < num_atoms; z++) {
                sum_exp += expf((v_row[z] + adv_a[z]) - max_logit);
            }

            /* Build CDF and extract quantiles */
            float cdf = 0.0f;
            float q10 = v_min;
            float q50 = v_min;
            float q90 = v_min;
            int got10 = 0, got50 = 0, got90 = 0;

            for (int z = 0; z < num_atoms; z++) {
                float prob = expf((v_row[z] + adv_a[z]) - max_logit) / sum_exp;
                cdf += prob;
                float z_val = v_min + (float)z * dz;
                if (!got10 && cdf >= 0.10f) { q10 = z_val; got10 = 1; }
                if (!got50 && cdf >= 0.50f) { q50 = z_val; got50 = 1; }
                if (!got90 && cdf >= 0.90f) { q90 = z_val; got90 = 1; }
            }

            /* Blend: exploring (alpha=0) → optimistic Q_90th,
             *        exploiting (alpha=1) → pessimistic Q_10th */
            float q_blended = (1.0f - alpha) * q90 + alpha * q10;

            q_select[(long long)i * total_actions + action_idx] = q_blended;
            action_idx++;
        }
        branch_logit_offset += N * n_d * num_atoms;
    }
}

• Step 3: Add the branch_graph_message_pass kernel. This implements a 4-node directed graph with gated message passing over the 12 Q-values (4 branches x 3 actions each):

/* ================================================================== */
/* Kernel: branch_graph_message_pass                                   */
/* ================================================================== */

/**
 * Cross-branch graph message passing with domain-specific edges.
 *
 * 4 directed edges encode branch dependencies:
 *   Edge 0: direction(0) -> magnitude(1)   — position sizing depends on direction
 *   Edge 1: direction(0) -> order(2)        — order type depends on market direction
 *   Edge 2: magnitude(1) -> urgency(3)      — position size affects timing pressure
 *   Edge 3: order(2)     -> urgency(3)      — limit vs market affects urgency
 *
 * Per edge: gate = sigmoid(W_gate @ [Q_src; Q_dst]), message = gate * (W_msg @ Q_src)
 * Per node: Q_new = Q_old + mean(incoming messages)
 *
 * Runs after cross_branch_q_attention, before diffusion refinement.
 * In-place residual update preserves original Q-value ordering.
 *
 * Grid: ceil(B/256), Block: 256. One thread per sample.
 *
 * @param q_values      [B, 12]   Q-values (in-place update)
 * @param graph_params  [60]      edge parameters: 4 edges x (W_gate[6] + W_msg[9]) = 60
 * @param B                       batch size
 */
extern "C" __global__ void branch_graph_message_pass(
    float* __restrict__ q_values,
    const float* __restrict__ graph_params,
    int B)
{
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= B) return;

    float* q = q_values + b * 12;

    /* Read current Q-values into local registers */
    float q_local[12];
    for (int i = 0; i < 12; i++) q_local[i] = q[i];

    /* Branch boundaries: dir=[0,3), mag=[3,6), ord=[6,9), urg=[9,12) */
    const int branch_start[4] = {0, 3, 6, 9};

    /* 4 edges: (src_branch, dst_branch) */
    const int edge_src[4] = {0, 0, 1, 2};
    const int edge_dst[4] = {1, 2, 3, 3};

    /* Accumulate incoming messages per node */
    float msg_accum[12] = {0.0f, 0.0f, 0.0f, 0.0f,
                           0.0f, 0.0f, 0.0f, 0.0f,
                           0.0f, 0.0f, 0.0f, 0.0f};
    int msg_count[4] = {0, 0, 0, 0};

    for (int e = 0; e < 4; e++) {
        int src_b = edge_src[e];
        int dst_b = edge_dst[e];
        int src_off = branch_start[src_b];
        int dst_off = branch_start[dst_b];

        /* W_gate[1, 6] at offset e * 15, W_msg[3, 3] at offset e * 15 + 6 */
        const float* W_gate = graph_params + e * 15;
        const float* W_msg  = graph_params + e * 15 + 6;

        /* gate = sigmoid(W_gate @ [Q_src; Q_dst]) */
        float gate_pre = 0.0f;
        for (int j = 0; j < 3; j++) gate_pre += W_gate[j]     * q_local[src_off + j];
        for (int j = 0; j < 3; j++) gate_pre += W_gate[3 + j] * q_local[dst_off + j];
        float gate = 1.0f / (1.0f + expf(-gate_pre));

        /* message = gate * (W_msg @ Q_src) */
        for (int i = 0; i < 3; i++) {
            float m = 0.0f;
            for (int j = 0; j < 3; j++) {
                m += W_msg[i * 3 + j] * q_local[src_off + j];
            }
            msg_accum[dst_off + i] += gate * m;
        }
        msg_count[dst_b]++;
    }

    /* Residual update: Q_new = Q_old + mean(incoming messages) */
    for (int d = 0; d < 4; d++) {
        if (msg_count[d] > 0) {
            float inv = 1.0f / (float)msg_count[d];
            int off = branch_start[d];
            for (int i = 0; i < 3; i++) {
                q[off + i] = q_local[off + i] + msg_accum[off + i] * inv;
            }
        }
    }
}

• Step 4: Add the kan_gate_combine kernel (replaces glu_combine). Evaluates cubic B-splines with learned coefficients plus residual connection, clamped to [0,1]:

/* ================================================================== */
/* Kernel: kan_gate_combine                                            */
/* ================================================================== */

/**
 * KAN Spline Activation gate — replaces sigmoid in GLU.
 *
 *   gate = clamp(sum_k coeff[k] * B_k(pre) + residual_w * pre, 0, 1)
 *   output = gate * value
 *
 * 8 cubic B-splines on uniform grid [-4, 4] per neuron.
 * Coefficients are learnable; basis functions are FIXED.
 * Initialized to approximate sigmoid (see param init step).
 *
 * Grid: ceil(n/256), Block: 256. One thread per element.
 *
 * @param output      [n]       output = gate * value
 * @param gate_pre    [n]       pre-activation input to gate (from W_gate @ input)
 * @param value       [n]       value path (from W_bdf @ input)
 * @param spline_coeff [AH, 8] spline coefficients per neuron (8 per neuron, tiled over batch)
 * @param residual_w  [AH]     residual weight per neuron (tiled over batch)
 * @param n                     total elements (batch_size * adv_h)
 * @param adv_h                 neuron count per sample (for tiling coeff/residual)
 */
extern "C" __global__ void kan_gate_combine(
    float* __restrict__ output,
    const float* __restrict__ gate_pre,
    const float* __restrict__ value,
    const float* __restrict__ spline_coeff,
    const float* __restrict__ residual_w,
    int n,
    int adv_h)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;

    float x = gate_pre[i];
    int neuron = i % adv_h;

    /* 8 cubic B-splines on uniform knots [-4, -3, -2, ..., 4] (9 knots, 8 spans)
     * Each B-spline B_k(x) is nonzero on [knot[k], knot[k+4]).
     * For uniform cubic: we use the de Boor recursive formula simplified for
     * uniform spacing (h=1.0).
     *
     * Simplified: find the span, compute the 4 nonzero basis values. */
    const float GRID_MIN = -4.0f;
    const float GRID_STEP = 1.0f;
    const int NUM_BASES = 8;

    /* Clamp x to grid range for basis evaluation */
    float x_clamped = fminf(fmaxf(x, GRID_MIN), GRID_MIN + NUM_BASES * GRID_STEP - 0.001f);

    /* Find span index: which interval [knot_i, knot_{i+1}) contains x */
    float t_norm = (x_clamped - GRID_MIN) / GRID_STEP;
    int span = (int)floorf(t_norm);
    span = min(max(span, 0), NUM_BASES - 1);
    float t = t_norm - (float)span; /* local parameter in [0, 1) */

    /* Cubic B-spline basis values (de Boor, uniform knots, h=1):
     * B_{span-1}(t) = (1-t)^3 / 6
     * B_{span}(t)   = (3t^3 - 6t^2 + 4) / 6
     * B_{span+1}(t) = (-3t^3 + 3t^2 + 3t + 1) / 6
     * B_{span+2}(t) = t^3 / 6 */
    float t2 = t * t;
    float t3 = t2 * t;
    float omt = 1.0f - t;
    float b0 = omt * omt * omt / 6.0f;
    float b1 = (3.0f * t3 - 6.0f * t2 + 4.0f) / 6.0f;
    float b2 = (-3.0f * t3 + 3.0f * t2 + 3.0f * t + 1.0f) / 6.0f;
    float b3 = t3 / 6.0f;

    /* Accumulate: gate_raw = sum(coeff[k] * B_k(x)) over the 4 nonzero bases */
    const float* c = spline_coeff + neuron * NUM_BASES;
    float gate_raw = 0.0f;

    int k0 = span - 1;
    if (k0 >= 0 && k0 < NUM_BASES) gate_raw += c[k0] * b0;
    if (span >= 0 && span < NUM_BASES) gate_raw += c[span] * b1;
    int k2 = span + 1;
    if (k2 >= 0 && k2 < NUM_BASES) gate_raw += c[k2] * b2;
    int k3 = span + 2;
    if (k3 >= 0 && k3 < NUM_BASES) gate_raw += c[k3] * b3;

    /* Add residual connection */
    gate_raw += residual_w[neuron] * x;

    /* Clamp to [0, 1] — gate must be a probability */
    float gate = fminf(fmaxf(gate_raw, 0.0f), 1.0f);

    output[i] = gate * value[i];
}

• Step 5: Add the kan_gate_backward kernel. Gradient through spline basis + residual + clamp:

/* ================================================================== */
/* Kernel: kan_gate_backward                                           */
/* ================================================================== */

/**
 * KAN Spline Activation backward pass.
 *
 * Computes gradients for:
 *   d_value[i]     = d_output[i] * gate
 *   d_gate_pre[i]  = d_output[i] * value[i] * d_gate/d_pre   (for upstream W_gate grad)
 *   d_coeff[k]    += d_output[i] * value[i] * d_gate/d_raw * B_k(pre)
 *   d_residual_w  += d_output[i] * value[i] * d_gate/d_raw * pre
 *
 * d_gate/d_raw = 1 if raw in (0,1), else 0 (clamp derivative).
 *
 * Gradient for coefficients uses atomicAdd (batch dimension reduction).
 *
 * Grid: ceil(n/256), Block: 256. One thread per element.
 */
extern "C" __global__ void kan_gate_backward(
    const float* __restrict__ d_output,
    const float* __restrict__ gate_pre,
    const float* __restrict__ value,
    const float* __restrict__ spline_coeff,
    const float* __restrict__ residual_w,
    float* __restrict__ d_gate_pre,
    float* __restrict__ d_value,
    float* __restrict__ d_spline_coeff,
    float* __restrict__ d_residual_w,
    int n,
    int adv_h)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;

    float x = gate_pre[i];
    int neuron = i % adv_h;

    const float GRID_MIN = -4.0f;
    const float GRID_STEP = 1.0f;
    const int NUM_BASES = 8;

    float x_clamped = fminf(fmaxf(x, GRID_MIN), GRID_MIN + NUM_BASES * GRID_STEP - 0.001f);
    float t_norm = (x_clamped - GRID_MIN) / GRID_STEP;
    int span = (int)floorf(t_norm);
    span = min(max(span, 0), NUM_BASES - 1);
    float t = t_norm - (float)span;

    float t2 = t * t;
    float t3 = t2 * t;
    float omt = 1.0f - t;
    float b0 = omt * omt * omt / 6.0f;
    float b1 = (3.0f * t3 - 6.0f * t2 + 4.0f) / 6.0f;
    float b2 = (-3.0f * t3 + 3.0f * t2 + 3.0f * t + 1.0f) / 6.0f;
    float b3 = t3 / 6.0f;

    /* Recompute gate_raw and gate */
    const float* c = spline_coeff + neuron * NUM_BASES;
    float gate_raw = 0.0f;
    int k0 = span - 1;
    if (k0 >= 0 && k0 < NUM_BASES) gate_raw += c[k0] * b0;
    if (span >= 0 && span < NUM_BASES) gate_raw += c[span] * b1;
    int k2 = span + 1;
    if (k2 >= 0 && k2 < NUM_BASES) gate_raw += c[k2] * b2;
    int k3_idx = span + 2;
    if (k3_idx >= 0 && k3_idx < NUM_BASES) gate_raw += c[k3_idx] * b3;
    gate_raw += residual_w[neuron] * x;
    float gate = fminf(fmaxf(gate_raw, 0.0f), 1.0f);

    /* Clamp derivative: d_gate/d_raw = 1 if raw in (0,1), else 0 */
    float clamp_mask = (gate_raw > 0.0f && gate_raw < 1.0f) ? 1.0f : 0.0f;

    float d_out = d_output[i];
    float val = value[i];

    /* d_value = d_output * gate */
    d_value[i] = d_out * gate;

    /* d_raw = d_output * value * clamp_mask */
    float d_raw = d_out * val * clamp_mask;

    /* d_gate_pre: chain through spline basis derivatives + residual */
    /* For simplicity, the upstream gradient for W_gate uses the sum of:
     * spline basis derivatives * coefficients + residual_w.
     * dB0/dt = -omt^2/2, dB1/dt = (3t^2 - 4t)/2, etc. (uniform cubic) */
    float db0_dt = -omt * omt / 2.0f;
    float db1_dt = (3.0f * t2 - 4.0f * t) / 2.0f;
    float db2_dt = (-3.0f * t2 + 2.0f * t + 0.5f);
    float db3_dt = t2 / 2.0f;

    float dgate_dx = 0.0f;
    if (k0 >= 0 && k0 < NUM_BASES) dgate_dx += c[k0] * db0_dt;
    if (span >= 0 && span < NUM_BASES) dgate_dx += c[span] * db1_dt;
    if (k2 >= 0 && k2 < NUM_BASES) dgate_dx += c[k2] * db2_dt;
    if (k3_idx >= 0 && k3_idx < NUM_BASES) dgate_dx += c[k3_idx] * db3_dt;
    /* Chain rule: dt/dx = 1/GRID_STEP = 1.0 for uniform grid */
    dgate_dx += residual_w[neuron];  /* residual contribution */

    d_gate_pre[i] = d_raw * dgate_dx;

    /* Accumulate d_coeff via atomicAdd (batch dimension reduction) */
    float* dc = d_spline_coeff + neuron * NUM_BASES;
    if (k0 >= 0 && k0 < NUM_BASES) atomicAdd(&dc[k0], d_raw * b0);
    if (span >= 0 && span < NUM_BASES) atomicAdd(&dc[span], d_raw * b1);
    if (k2 >= 0 && k2 < NUM_BASES) atomicAdd(&dc[k2], d_raw * b2);
    if (k3_idx >= 0 && k3_idx < NUM_BASES) atomicAdd(&dc[k3_idx], d_raw * b3);

    /* Accumulate d_residual_w via atomicAdd */
    atomicAdd(&d_residual_w + neuron, d_raw * x);
}

• Step 6: Add the q_denoise_step kernel. One denoising step: FC-SiLU-FC residual block conditioned on distributional variance:

/* ================================================================== */
/* Kernel: q_denoise_step                                              */
/* ================================================================== */

/**
 * Single diffusion denoising step for Q-value refinement.
 *
 *   noise_level = sqrt(var) * (1 - step_k / K)
 *   input = [Q_prev; noise_level]              // [B, 24]
 *   h = SiLU(W1 @ input + b1)                  // [B, 24]
 *   residual = W2 @ h + b2                      // [B, 12]
 *   Q_next = Q_prev + 0.1 * residual           // small step
 *
 * SiLU(x) = x * sigmoid(x).
 *
 * Grid: ceil(B/256), Block: 256. One thread per sample.
 *
 * @param q_prev       [B, 12]   Q-values from previous step (read)
 * @param q_next       [B, 12]   Q-values after this step (write; may alias q_prev)
 * @param q_variance   [B, 12]   Var[Q] per action from C51 distribution
 * @param W1           [24, 24]  first FC weight
 * @param b1           [24]      first FC bias
 * @param W2           [12, 24]  second FC weight
 * @param b2           [12]      second FC bias
 * @param B                      batch size
 * @param step_k                 current step (1, 2, or 3)
 * @param total_steps            total steps K (3)
 */
extern "C" __global__ void q_denoise_step(
    const float* __restrict__ q_prev,
    float* __restrict__ q_next,
    const float* __restrict__ q_variance,
    const float* __restrict__ W1,
    const float* __restrict__ b1,
    const float* __restrict__ W2,
    const float* __restrict__ b2,
    int B,
    int step_k,
    int total_steps)
{
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= B) return;

    const int D = 12;
    const int H = 24; /* hidden = 2*D */

    const float* q_in = q_prev + b * D;
    const float* var_in = q_variance + b * D;

    /* Build input[24] = [Q_prev[12]; noise_level[12]] */
    float input[24];
    float schedule = 1.0f - (float)step_k / (float)total_steps;
    for (int i = 0; i < D; i++) {
        input[i] = q_in[i];
        input[D + i] = sqrtf(fmaxf(var_in[i], 0.0f)) * schedule;
    }

    /* h = SiLU(W1 @ input + b1) */
    float h[24];
    for (int i = 0; i < H; i++) {
        float acc = b1[i];
        for (int j = 0; j < H; j++) {
            acc += W1[i * H + j] * input[j];
        }
        /* SiLU: x * sigmoid(x) */
        float sig = 1.0f / (1.0f + expf(-acc));
        h[i] = acc * sig;
    }

    /* residual = W2 @ h + b2 */
    float* q_out = q_next + b * D;
    for (int i = 0; i < D; i++) {
        float acc = b2[i];
        for (int j = 0; j < H; j++) {
            acc += W2[i * H + j] * h[j];
        }
        /* Q_next = Q_prev + 0.1 * residual */
        q_out[i] = q_in[i] + 0.1f * acc;
    }
}

• Step 7: Add the concat_ofi_features kernel. Copies vsn_masked output and appends 3 OFI features for order/urgency branches:

/* ================================================================== */
/* Kernel: concat_ofi_features                                         */
/* ================================================================== */

/**
 * Concatenate VSN-masked trunk output with OFI microstructure features.
 *
 * For order branch (d=2):  out[b] = [vsn_masked[b, :SH2]; ofi_order[b, 3]]
 * For urgency branch (d=3): out[b] = [vsn_masked[b, :SH2]; ofi_urgency[b, 3]]
 *
 * OFI feature indices (from 50-dim feature vector, OFI at dims 42-49):
 *   order_extra:   [42=bid_ask_spread, 43=depth_imbalance_L1, 44=queue_pressure]
 *   urgency_extra: [45=spread_velocity, 46=depth_change_rate, 47=trade_arrival_rate]
 *
 * Grid: ceil((B * (SH2+3)) / 256), Block: 256.
 *
 * @param output        [B, SH2+3]  output concat buffer
 * @param vsn_masked    [B, SH2]    VSN-masked trunk output
 * @param features      [B, SD]     raw feature vector (contains OFI at fixed indices)
 * @param B                         batch size
 * @param SH2                       shared_h2 dimension
 * @param SD                        state_dim (feature vector width)
 * @param ofi_start_idx             first OFI index for this branch (42 for order, 45 for urgency)
 */
extern "C" __global__ void concat_ofi_features(
    float* __restrict__ output,
    const float* __restrict__ vsn_masked,
    const float* __restrict__ features,
    int B,
    int SH2,
    int SD,
    int ofi_start_idx)
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * (SH2 + 3);
    if (idx >= total) return;

    int b = idx / (SH2 + 3);
    int col = idx % (SH2 + 3);

    if (col < SH2) {
        output[idx] = vsn_masked[b * SH2 + col];
    } else {
        int ofi_col = col - SH2; /* 0, 1, or 2 */
        output[idx] = features[b * SD + ofi_start_idx + ofi_col];
    }
}

• Step 8: Verify all kernels compile by running:

SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -20

The build.rs compiles experience_kernels.cu to a cubin. If nvcc is not available locally, check for the cubin pattern in the build output. No runtime test at this stage.

• Step 9: Commit:

git add crates/ml/src/cuda_pipeline/experience_kernels.cu
git commit -m "feat: add 7 CUDA kernels for supervised architecture transfer

Add quantile_q_select, branch_graph_message_pass, kan_gate_combine,
kan_gate_backward, q_denoise_step, concat_ofi_features kernels to
experience_kernels.cu. All compile as part of existing cubin pipeline.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>"

---

Task 2: Quantile Q-Select Integration

Files:

• crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
• crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
• crates/ml/src/trainers/dqn/fused_training.rs

Wire the quantile_q_select kernel into experience collection to replace Boltzmann-on-E[Q] action selection. Zero new parameters -- uses existing C51 atoms.

• Step 1: Read crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs around line 2794 where kernels are loaded from cpbi_module. Add loading of the new kernel:

In gpu_dqn_trainer.rs, after line ~2806 (let glu_backward_kernel = cpbi_module.load_function("glu_backward")...), add:

        let quantile_q_select_kernel = cpbi_module.load_function("quantile_q_select")
            .map_err(|e| MLError::ModelError(format!("quantile_q_select load: {e}")))?;

• Step 2: Add the kernel function field to FusedTrainingCtx (in gpu_dqn_trainer.rs). Find the struct fields section (around line 661 where glu_combine_kernel is stored) and add:

    quantile_q_select_kernel: CudaFunction,

Store it in the constructor where the struct is built (around line 3346).

• Step 3: Add a q_select_buf allocation in the constructor (around line 2717 where q_coord_buf is allocated):

        let q_select_buf = alloc_f32(&stream, b * 12, "q_select_buf")?;

Add the field to the struct and a q_select_buf_ptr() accessor:

    pub(crate) fn q_select_buf_ptr(&self) -> u64 {
        self.q_select_buf.raw_ptr()
    }

• Step 4: Add launch_quantile_q_select method to GpuDqnTrainer:

    /// Launch quantile Q-select kernel: C51 atoms → quantile-blended Q-values.
    /// Writes to q_select_buf [batch_size, 12]. Uses iqn_readiness for blend alpha.
    pub(crate) fn launch_quantile_q_select(&self, batch_size: usize) -> Result<(), MLError> {
        let blocks = ((batch_size as u32 + 255) / 256).max(1);
        let n = batch_size as i32;
        let na = self.config.num_atoms as i32;
        let b0 = self.config.branch_0_size as i32;
        let b1 = self.config.branch_1_size as i32;
        let b2 = self.config.branch_2_size as i32;
        let b3 = self.config.branch_3_size as i32;
        let on_v_ptr = self.ptrs.on_v_logits_buf;
        let on_b_ptr = self.ptrs.on_b_logits_buf;
        let q_select_ptr = self.q_select_buf.raw_ptr();
        let support_ptr = self.per_sample_support_ptr;
        let readiness = self.iqn_readiness;
        unsafe {
            self.stream
                .launch_builder(&self.quantile_q_select_kernel)
                .arg(&on_v_ptr)
                .arg(&on_b_ptr)
                .arg(&q_select_ptr)
                .arg(&n)
                .arg(&na)
                .arg(&b0)
                .arg(&b1)
                .arg(&b2)
                .arg(&b3)
                .arg(&support_ptr)
                .arg(&readiness)
                .launch(LaunchConfig {
                    grid_dim: (blocks, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("quantile_q_select: {e}")))?;
        }
        Ok(())
    }

• Step 5: Add passthrough in fused_training.rs. Find the launch_q_attention passthrough (around line 1981) and add after it:

    /// Launch quantile Q-select: C51 atoms → quantile-blended Q for action selection.
    pub(crate) fn launch_quantile_q_select(&self, batch_size: usize) -> Result<()> {
        self.trainer.launch_quantile_q_select(batch_size)
            .map_err(|e| anyhow::anyhow!("launch_quantile_q_select: {e}"))
    }

    /// Raw pointer to quantile Q-select output buffer [batch_size, 12].
    pub(crate) fn q_select_buf_ptr(&self) -> u64 {
        self.trainer.q_select_buf_ptr()
    }

• Step 6: Wire into experience collection. In gpu_experience_collector.rs, find the compute_expected_q launch in the timestep loop (around line 1981). After compute_expected_q runs, the collector currently passes q_values (E[Q]) to experience_action_select. Instead, call quantile_q_select and pass the result to action select.

The collector needs a reference to the quantile kernel. Add a quantile_q_select_kernel: Option<CudaFunction> field and a q_select_buf: Option<CudaSlice<f32>> to GpuExperienceCollector. Load the kernel from the same module used for compute_expected_q (the experience_kernels cubin loaded at line ~2586).

In the timestep loop, after compute_expected_q populates q_values, add:

// Quantile Q-select: overwrite q_values with quantile-blended Q for action selection
if let (Some(ref qqs_kernel), Some(ref qsel_buf)) = (&self.quantile_q_select_kernel, &self.q_select_buf) {
    let blocks = ((n_episodes as u32 + 255) / 256).max(1);
    let na = self.num_atoms as i32;
    let support_ptr = self.per_sample_support_ptr;
    let readiness = self.iqn_readiness;
    unsafe {
        self.stream
            .launch_builder(qqs_kernel)
            .arg(&v_logits_ptr)
            .arg(&b_logits_ptr)
            .arg(&qsel_buf.raw_ptr())
            .arg(&n_episodes)
            .arg(&na)
            .arg(&b0)
            .arg(&b1)
            .arg(&b2)
            .arg(&b3)
            .arg(&support_ptr)
            .arg(&readiness)
            .launch(LaunchConfig {
                grid_dim: (blocks, 1, 1),
                block_dim: (256, 1, 1),
                shared_mem_bytes: 0,
            })
            .map_err(|e| MLError::ModelError(format!("quantile_q_select t={t}: {e}")))?;
    }
    // Use q_select_buf instead of q_values for action selection
    q_values_ptr_for_action = qsel_buf.raw_ptr();
}

• Step 7: Add iqn_readiness field to GpuExperienceCollector and a setter method:

    pub fn set_iqn_readiness(&mut self, readiness: f32) {
        self.iqn_readiness = readiness;
    }

Wire the setter call in the training loop (in fused_training.rs or training_loop.rs) after update_iqn_readiness is called, propagating the value to the collector.

• Step 8: Build check:

SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -10

• Step 9: Commit:

git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs crates/ml/src/cuda_pipeline/gpu_experience_collector.rs crates/ml/src/trainers/dqn/fused_training.rs
git commit -m "feat: wire quantile Q-select into experience collection action selection

Replace Boltzmann-on-E[Q] with distributional quantile blend:
Q_select = (1-alpha)*Q_90th + alpha*Q_10th where alpha = iqn_readiness.
Zero new parameters — reuses existing C51 atom distributions.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>"

---

Task 3: Graph Message Passing Integration

Files:

• crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
• crates/ml/src/trainers/dqn/fused_training.rs

Wire the branch_graph_message_pass kernel after Q-attention. 60 new parameters stored in a separate GPU buffer with its own Adam optimizer (same pattern as q_attn_params).

• Step 1: In gpu_dqn_trainer.rs, add graph params allocation in the constructor (around line 2711 where q_attn_params is allocated):

        let graph_params = stream.alloc_zeros::<f32>(60)
            .map_err(|e| MLError::ModelError(format!("alloc graph_params: {e}")))?;
        // Initialize graph_params: W_gate to small random, W_msg to identity-like
        {
            let mut init = vec![0.0_f32; 60];
            for e in 0..4 {
                // W_msg [3,3] initialized to 0.1 * I (small identity — passes through weakly)
                let msg_off = e * 15 + 6;
                init[msg_off + 0] = 0.1; // [0,0]
                init[msg_off + 4] = 0.1; // [1,1]
                init[msg_off + 8] = 0.1; // [2,2]
            }
            stream.memcpy_htod(&graph_params, &init)
                .map_err(|e| MLError::ModelError(format!("graph_params init: {e}")))?;
        }

• Step 2: Add Adam state buffers for graph params (same pattern as selectivity Adam):

        let graph_m = stream.alloc_zeros::<f32>(60)
            .map_err(|e| MLError::ModelError(format!("alloc graph_m: {e}")))?;
        let graph_v = stream.alloc_zeros::<f32>(60)
            .map_err(|e| MLError::ModelError(format!("alloc graph_v: {e}")))?;
        let graph_grad = stream.alloc_zeros::<f32>(60)
            .map_err(|e| MLError::ModelError(format!("alloc graph_grad: {e}")))?;

• Step 3: Add struct fields and store them:

    graph_params: CudaSlice<f32>,      // [60]
    graph_m: CudaSlice<f32>,           // [60]
    graph_v: CudaSlice<f32>,           // [60]
    graph_grad: CudaSlice<f32>,        // [60]
    graph_kernel: CudaFunction,
    graph_adam_step: i32,

• Step 4: Load the kernel from the cpbi_module (after quantile_q_select_kernel):

        let graph_kernel = cpbi_module.load_function("branch_graph_message_pass")
            .map_err(|e| MLError::ModelError(format!("branch_graph_message_pass load: {e}")))?;

• Step 5: Add launch_graph_message_pass method to GpuDqnTrainer:

    /// Launch graph message passing on q_coord_buf (in-place update).
    /// Must be called after launch_q_attention.
    pub(crate) fn launch_graph_message_pass(&self, batch_size: usize) -> Result<(), MLError> {
        let blocks = ((batch_size as u32 + 255) / 256).max(1);
        let b = batch_size as i32;
        let q_coord_ptr = self.q_coord_buf.raw_ptr();
        let graph_params_ptr = self.graph_params.raw_ptr();
        unsafe {
            self.stream
                .launch_builder(&self.graph_kernel)
                .arg(&q_coord_ptr)
                .arg(&graph_params_ptr)
                .arg(&b)
                .launch(LaunchConfig {
                    grid_dim: (blocks, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("branch_graph_message_pass: {e}")))?;
        }
        Ok(())
    }

• Step 6: Wire into the Q-value pipeline. In reduce_current_q_stats (around line 4758), after launch_q_attention(batch_size)?;, add:

        // Graph message passing: q_coord_buf → q_coord_buf (in-place residual).
        self.launch_graph_message_pass(batch_size)?;

• Step 7: Add passthrough in fused_training.rs:

    pub(crate) fn launch_graph_message_pass(&self, batch_size: usize) -> Result<()> {
        self.trainer.launch_graph_message_pass(batch_size)
            .map_err(|e| anyhow::anyhow!("launch_graph_message_pass: {e}"))
    }

• Step 8: Build check:

SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -10

• Step 9: Commit:

git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs crates/ml/src/trainers/dqn/fused_training.rs
git commit -m "feat: wire graph message passing after Q-attention

Add branch_graph_message_pass with 4 domain edges (dir->mag, dir->ord,
mag->urg, ord->urg). 60 params with separate Adam. Runs in-place on
q_coord_buf after cross_branch_q_attention.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>"

---

Task 4: KAN Spline Gates -- Replace GLU

Files:

• crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
• crates/ml/src/cuda_pipeline/batched_forward.rs
• crates/ml/src/cuda_pipeline/batched_backward.rs

Replace glu_combine sigmoid gating with learned B-spline activations. 9,216 new parameters: 4 branches x 256 neurons x (8 spline coefficients + 1 residual weight).

• Step 1: Increase NUM_WEIGHT_TENSORS from 42 to 50. In gpu_dqn_trainer.rs line 355:

Change:

pub(crate) const NUM_WEIGHT_TENSORS: usize = 42;

To:

pub(crate) const NUM_WEIGHT_TENSORS: usize = 50;

• Step 2: Extend compute_param_sizes to include 8 new tensors (indices 42-49). After the existing index 41 entry, add:

        // ── KAN spline coefficients + residual weights ──
        cfg.adv_h * 8,                                       // [42] spline_coeff_0 [AH, 8]
        cfg.adv_h,                                            // [43] residual_w_0   [AH]
        cfg.adv_h * 8,                                       // [44] spline_coeff_1 [AH, 8]
        cfg.adv_h,                                            // [45] residual_w_1   [AH]
        cfg.adv_h * 8,                                       // [46] spline_coeff_2 [AH, 8]
        cfg.adv_h,                                            // [47] residual_w_2   [AH]
        cfg.adv_h * 8,                                       // [48] spline_coeff_3 [AH, 8]
        cfg.adv_h,                                            // [49] residual_w_3   [AH]

• Step 3: Update the doc comment above compute_param_sizes to reflect the new range:

/// Tensors 42-49: KAN spline coefficients and residual weights per branch.

• Step 4: Initialize KAN spline coefficients to approximate sigmoid. In the Xavier init section (around line 6445 where fan_dims is defined), extend the fan_dims array to 50 entries. For the 8 new entries, use:

            (cfg.adv_h * 8, 1),  // [42] spline_coeff_0 (not truly a weight matrix — use custom init)
            (cfg.adv_h, 1),       // [43] residual_w_0
            (cfg.adv_h * 8, 1),  // [44] spline_coeff_1
            (cfg.adv_h, 1),       // [45] residual_w_1
            (cfg.adv_h * 8, 1),  // [46] spline_coeff_2
            (cfg.adv_h, 1),       // [47] residual_w_2
            (cfg.adv_h * 8, 1),  // [48] spline_coeff_3
            (cfg.adv_h, 1),       // [49] residual_w_3

Then add custom initialization after Xavier init to set spline coefficients that approximate sigmoid and residual_w to 0:

        // KAN spline init: approximate sigmoid(x) on [-4, 4] with B-spline basis.
        // Precomputed coefficients that give sigmoid-like gate output.
        // Evaluated at grid midpoints [-3.5, -2.5, -1.5, -0.5, 0.5, 1.5, 2.5, 3.5]:
        let sigmoid_approx_coeffs: [f32; 8] = [
            0.03, 0.07, 0.18, 0.38, 0.62, 0.82, 0.93, 0.97
        ];
        for branch in 0..4 {
            let coeff_idx = 42 + branch * 2;     // spline_coeff index
            let resid_idx = 42 + branch * 2 + 1; // residual_w index
            let coeff_offset = padded_byte_offset(&sizes, coeff_idx) / 4;
            let resid_offset = padded_byte_offset(&sizes, resid_idx) / 4;
            for neuron in 0..cfg.adv_h {
                for k in 0..8 {
                    host_params[coeff_offset + neuron * 8 + k] = sigmoid_approx_coeffs[k];
                }
                host_params[resid_offset + neuron] = 0.0; // residual starts at zero
            }
        }

• Step 5: Load kan_gate_combine and kan_gate_backward kernels from the cpbi_module:

        let kan_gate_combine_kernel = cpbi_module.load_function("kan_gate_combine")
            .map_err(|e| MLError::ModelError(format!("kan_gate_combine load: {e}")))?;
        let kan_gate_backward_kernel = cpbi_module.load_function("kan_gate_backward")
            .map_err(|e| MLError::ModelError(format!("kan_gate_backward load: {e}")))?;

Add fields to the struct:

    kan_gate_combine_kernel: CudaFunction,
    kan_gate_backward_kernel: CudaFunction,

• Step 6: Allocate KAN gradient scratch buffers (for the atomicAdd reduction in backward):

        // KAN backward scratch: d_spline_coeff [AH, 8] and d_residual_w [AH] per branch
        let kan_d_spline_coeff = [
            alloc_f32(&stream, config.adv_h * 8, "kan_dsc0")?,
            alloc_f32(&stream, config.adv_h * 8, "kan_dsc1")?,
            alloc_f32(&stream, config.adv_h * 8, "kan_dsc2")?,
            alloc_f32(&stream, config.adv_h * 8, "kan_dsc3")?,
        ];
        let kan_d_residual_w = [
            alloc_f32(&stream, config.adv_h, "kan_drw0")?,
            alloc_f32(&stream, config.adv_h, "kan_drw1")?,
            alloc_f32(&stream, config.adv_h, "kan_drw2")?,
            alloc_f32(&stream, config.adv_h, "kan_drw3")?,
        ];

• Step 7: Replace glu_combine kernel calls with kan_gate_combine. In batched_forward.rs, find the forward_online_raw method where glu_combine is launched (around line 958). Replace:

        // OLD: glu_combine(output, gate_pre, value, n)
        // NEW: kan_gate_combine(output, gate_pre, value, spline_coeff, residual_w, n, adv_h)

The key change: where the forward pass currently calls glu_k (the glu_combine_kernel), instead call the kan_gate_combine_kernel with the additional spline_coeff and residual_w pointers. The spline_coeff pointer is computed from params_buf using padded_byte_offset(&param_sizes, 42 + d * 2), and residual_w from padded_byte_offset(&param_sizes, 43 + d * 2).

Wire the kan_gate_combine_kernel through CublasGemmSet::wire_vsn_glu (or rename to wire_vsn_kan). The forward pass needs access to params_buf raw pointer + param_sizes to compute offsets. Pass these as additional arguments to wire_vsn_glu:

    pub fn wire_vsn_kan(
        &mut self,
        // ... existing args ...
        kan_gate_combine_kernel: CudaFunction,
        params_buf_ptr: u64,
        param_sizes: &[usize; NUM_WEIGHT_TENSORS],
    ) {
        self.kan_gate_combine_kernel = Some(kan_gate_combine_kernel);
        self.kan_params_buf_ptr = params_buf_ptr;
        // Pre-compute spline_coeff and residual_w offsets per branch
        for d in 0..4 {
            self.kan_spline_coeff_offset[d] = padded_byte_offset(param_sizes, 42 + d * 2);
            self.kan_residual_w_offset[d] = padded_byte_offset(param_sizes, 43 + d * 2);
        }
    }

• Step 8: Replace glu_backward kernel calls with kan_gate_backward in batched_backward.rs. The backward pass currently calls glu_backward_kernel to split d_h_bd into d_value and d_gate_pre. Replace with kan_gate_backward which additionally produces d_spline_coeff and d_residual_w gradients. After the backward pass, accumulate these into grad_buf at the appropriate offsets (indices 42-49).

• Step 9: Build check:

SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -10

• Step 10: Commit:

git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs crates/ml/src/cuda_pipeline/batched_forward.rs crates/ml/src/cuda_pipeline/batched_backward.rs
git commit -m "feat: replace GLU sigmoid gates with KAN B-spline activations

NUM_WEIGHT_TENSORS 42->50. New indices 42-49: per-branch spline_coeff[AH,8]
and residual_w[AH]. 9216 total new params. Coefficients initialized to
approximate sigmoid. kan_gate_combine replaces glu_combine in forward;
kan_gate_backward replaces glu_backward with atomicAdd coeff gradients.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>"

---

Task 5: Diffusion Q-Refinement

Files:

• crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
• crates/ml/src/trainers/dqn/fused_training.rs

Add 3-step iterative Q-value denoising conditioned on distributional variance. 2,772 new parameters with separate Adam optimizer.

• Step 1: Add diffusion parameter buffers in the GpuDqnTrainer constructor. Each step has W1[24,24]+b1[24]+W2[12,24]+b2[12] = 924 params. 3 steps = 2,772:

        // Diffusion Q-refinement: 3 denoiser steps, each with W1[24,24]+b1[24]+W2[12,24]+b2[12]
        let denoise_params_per_step = 24 * 24 + 24 + 12 * 24 + 12; // = 924
        let denoise_total_params = denoise_params_per_step * 3;     // = 2772
        let denoise_params = stream.alloc_zeros::<f32>(denoise_total_params)
            .map_err(|e| MLError::ModelError(format!("alloc denoise_params: {e}")))?;
        let denoise_m = stream.alloc_zeros::<f32>(denoise_total_params)
            .map_err(|e| MLError::ModelError(format!("alloc denoise_m: {e}")))?;
        let denoise_v = stream.alloc_zeros::<f32>(denoise_total_params)
            .map_err(|e| MLError::ModelError(format!("alloc denoise_v: {e}")))?;
        let denoise_grad = stream.alloc_zeros::<f32>(denoise_total_params)
            .map_err(|e| MLError::ModelError(format!("alloc denoise_grad: {e}")))?;
        // Initialize W1, W2 with small random (Xavier), biases zero
        {
            let mut init = vec![0.0_f32; denoise_total_params];
            let mut rng = 0xDEADBEEFu64;
            for step in 0..3 {
                let base = step * denoise_params_per_step;
                // W1[24,24]: Xavier scale = sqrt(2/(24+24)) = 0.204
                for i in 0..(24 * 24) {
                    rng = rng.wrapping_mul(6364136223846793005).wrapping_add(1);
                    let u = (rng >> 33) as f32 / (1u64 << 31) as f32 - 0.5;
                    init[base + i] = u * 0.408; // 2 * 0.204
                }
                // b1[24] = 0 (already zero)
                // W2[12,24]: Xavier scale = sqrt(2/(12+24)) = 0.236
                let w2_off = base + 24 * 24 + 24;
                for i in 0..(12 * 24) {
                    rng = rng.wrapping_mul(6364136223846793005).wrapping_add(1);
                    let u = (rng >> 33) as f32 / (1u64 << 31) as f32 - 0.5;
                    init[w2_off + i] = u * 0.472;
                }
                // b2[12] = 0 (already zero)
            }
            stream.memcpy_htod(&denoise_params, &init)
                .map_err(|e| MLError::ModelError(format!("denoise_params init: {e}")))?;
        }

• Step 2: Add q_variance_buf allocation for storing Var[Q] output from compute_expected_q:

        let q_variance_buf = alloc_f32(&stream, b * 12, "q_variance_buf")?;

Currently compute_expected_q passes null_q_var = 0u64 for the variance output. Change the Q-stats path to pass q_variance_buf.raw_ptr() instead, so Var[Q] is populated.

• Step 3: Add a q_denoise_scratch buffer for intermediate Q-values during the 3 steps:

        let q_denoise_scratch = alloc_f32(&stream, b * 12, "q_denoise_scratch")?;

• Step 4: Load the kernel:

        let q_denoise_kernel = cpbi_module.load_function("q_denoise_step")
            .map_err(|e| MLError::ModelError(format!("q_denoise_step load: {e}")))?;

• Step 5: Add struct fields:

    denoise_params: CudaSlice<f32>,      // [2772]
    denoise_m: CudaSlice<f32>,
    denoise_v: CudaSlice<f32>,
    denoise_grad: CudaSlice<f32>,
    denoise_adam_step: i32,
    q_variance_buf: CudaSlice<f32>,      // [B, 12]
    q_denoise_scratch: CudaSlice<f32>,   // [B, 12]
    q_denoise_kernel: CudaFunction,

• Step 6: Add launch_q_denoise method:

    /// Run 3-step diffusion Q-refinement on q_coord_buf (in-place).
    /// Requires q_variance_buf to be populated (from compute_expected_q).
    pub(crate) fn launch_q_denoise(&self, batch_size: usize) -> Result<(), MLError> {
        let blocks = ((batch_size as u32 + 255) / 256).max(1);
        let b = batch_size as i32;
        let total_steps = 3i32;
        let q_var_ptr = self.q_variance_buf.raw_ptr();
        let params_per_step = (24 * 24 + 24 + 12 * 24 + 12) as usize;

        // Ping-pong between q_coord_buf and q_denoise_scratch
        let buf_a = self.q_coord_buf.raw_ptr();
        let buf_b = self.q_denoise_scratch.raw_ptr();

        for k in 1..=3i32 {
            let (src, dst) = if k % 2 == 1 { (buf_a, buf_b) } else { (buf_b, buf_a) };
            let step_offset = ((k - 1) as usize) * params_per_step;
            let w1_ptr = self.denoise_params.raw_ptr() + (step_offset * 4) as u64;
            let b1_ptr = w1_ptr + (24 * 24 * 4) as u64;
            let w2_ptr = b1_ptr + (24 * 4) as u64;
            let b2_ptr = w2_ptr + (12 * 24 * 4) as u64;

            unsafe {
                self.stream
                    .launch_builder(&self.q_denoise_kernel)
                    .arg(&src)
                    .arg(&dst)
                    .arg(&q_var_ptr)
                    .arg(&w1_ptr)
                    .arg(&b1_ptr)
                    .arg(&w2_ptr)
                    .arg(&b2_ptr)
                    .arg(&b)
                    .arg(&k)
                    .arg(&total_steps)
                    .launch(LaunchConfig {
                        grid_dim: (blocks, 1, 1),
                        block_dim: (256, 1, 1),
                        shared_mem_bytes: 0,
                    })
                    .map_err(|e| MLError::ModelError(format!("q_denoise_step k={k}: {e}")))?;
            }
        }

        // After 3 steps (odd count), result is in buf_b. Copy back to q_coord_buf.
        unsafe {
            cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
                buf_a, buf_b,
                (batch_size * 12 * 4) as u64,
                self.stream.cu_stream(),
            );
        }
        Ok(())
    }

• Step 7: Wire into the Q-value pipeline. In reduce_current_q_stats, after graph message passing, before the stats reduction:

        // Diffusion Q-refinement: 3 denoising steps on q_coord_buf.
        self.launch_q_denoise(batch_size)?;

• Step 8: Enable q_variance output in reduce_current_q_stats. Change the null_q_var to q_variance_buf.raw_ptr():

        let q_var_ptr = self.q_variance_buf.raw_ptr();
        // In the compute_expected_q launch, replace:
        //   .arg(&null_q_var)
        // with:
        //   .arg(&q_var_ptr)

• Step 9: Add passthrough in fused_training.rs:

    pub(crate) fn launch_q_denoise(&self, batch_size: usize) -> Result<()> {
        self.trainer.launch_q_denoise(batch_size)
            .map_err(|e| anyhow::anyhow!("launch_q_denoise: {e}"))
    }

• Step 10: Build check:

SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -10

• Step 11: Commit:

git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs crates/ml/src/trainers/dqn/fused_training.rs
git commit -m "feat: add 3-step diffusion Q-refinement after graph message passing

2772 params (3 x FC-SiLU-FC blocks) with separate Adam. Conditioned on
Var[Q] from C51 distribution. Iteratively refines q_coord_buf with 0.1x
residual steps. Runs after graph_message_pass, before Q-stats reduction.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>"

---

Task 6: TLOB Microstructure Injection

Files:

• crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
• crates/ml/src/cuda_pipeline/batched_forward.rs
• crates/ml/src/cuda_pipeline/batched_backward.rs

Widen order branch (d=2) and urgency branch (d=3) FC input from SH2 to SH2+3 by concatenating OFI microstructure features. Same pattern as the existing direction-conditioned magnitude branch (d=1).

• Step 1: Change param_sizes indices [16], [20], [38], [40] to include +3 input dimension. In compute_param_sizes:

Change:

        cfg.adv_h * cfg.shared_h2,                          // [16] w_b2fc

To:

        cfg.adv_h * (cfg.shared_h2 + 3),                    // [16] w_b2fc (TLOB: SH2+3 input)

Change:

        cfg.adv_h * cfg.shared_h2,                          // [20] w_b3fc

To:

        cfg.adv_h * (cfg.shared_h2 + 3),                    // [20] w_b3fc (TLOB: SH2+3 input)

Change:

        cfg.adv_h * cfg.shared_h2,                           // [38] w_gate_2

To:

        cfg.adv_h * (cfg.shared_h2 + 3),                    // [38] w_gate_2 (TLOB: SH2+3 input)

Change:

        cfg.adv_h * cfg.shared_h2,                           // [40] w_gate_3

To:

        cfg.adv_h * (cfg.shared_h2 + 3),                    // [40] w_gate_3 (TLOB: SH2+3 input)

• Step 2: Add ord_concat_dim and urg_concat_dim fields to CublasGemmSet (in batched_forward.rs), analogous to the existing mag_concat_dim:

    /// Order branch (d==2) FC input dimension: shared_h2 + 3
    ord_concat_dim: usize,
    /// Urgency branch (d==3) FC input dimension: shared_h2 + 3
    urg_concat_dim: usize,

Initialize both to shared_h2 + 3 in the constructor.

• Step 3: Add concat buffer allocations in GpuDqnTrainer constructor:

        let ord_concat_buf = alloc_f32(&stream, b * (config.shared_h2 + 3), "ord_concat_buf")?;
        let urg_concat_buf = alloc_f32(&stream, b * (config.shared_h2 + 3), "urg_concat_buf")?;

• Step 4: Load the concat_ofi_features kernel:

        let concat_ofi_kernel = cpbi_module.load_function("concat_ofi_features")
            .map_err(|e| MLError::ModelError(format!("concat_ofi_features load: {e}")))?;

• Step 5: Add launch_ofi_concat method:

    /// Concatenate VSN-masked output with OFI features for order (d=2) and urgency (d=3) branches.
    pub(crate) fn launch_ofi_concat(&self, batch_size: usize) -> Result<(), MLError> {
        let sh2 = self.config.shared_h2 as i32;
        let sd = self.config.state_dim as i32;
        let b = batch_size as i32;
        let total_ord = b * (sh2 + 3);
        let total_urg = b * (sh2 + 3);
        let blocks_ord = ((total_ord as u32 + 255) / 256).max(1);
        let blocks_urg = ((total_urg as u32 + 255) / 256).max(1);

        let vsn_ptr = self.vsn_masked_buf.raw_ptr();
        let states_ptr = self.ptrs.states_buf;
        let ord_ptr = self.ord_concat_buf.raw_ptr();
        let urg_ptr = self.urg_concat_buf.raw_ptr();

        // Order branch: OFI features at indices 42, 43, 44
        let ofi_ord_start = 42i32;
        unsafe {
            self.stream
                .launch_builder(&self.concat_ofi_kernel)
                .arg(&ord_ptr)
                .arg(&vsn_ptr)
                .arg(&states_ptr)
                .arg(&b)
                .arg(&sh2)
                .arg(&sd)
                .arg(&ofi_ord_start)
                .launch(LaunchConfig {
                    grid_dim: (blocks_ord, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("concat_ofi order: {e}")))?;
        }

        // Urgency branch: OFI features at indices 45, 46, 47
        let ofi_urg_start = 45i32;
        unsafe {
            self.stream
                .launch_builder(&self.concat_ofi_kernel)
                .arg(&urg_ptr)
                .arg(&vsn_ptr)
                .arg(&states_ptr)
                .arg(&b)
                .arg(&sh2)
                .arg(&sd)
                .arg(&ofi_urg_start)
                .launch(LaunchConfig {
                    grid_dim: (blocks_urg, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("concat_ofi urgency: {e}")))?;
        }
        Ok(())
    }

• Step 6: Modify the forward pass in batched_forward.rs to use the concat buffers for branches 2 and 3. Follow the same pattern as the existing mag_concat_ptr for branch 1. Add ord_concat_ptr and urg_concat_ptr parameters to forward_online_raw:

In the branch loop, for d == 2:

let (fc_input, fc_k) = if d == 2 && ord_concat_ptr != 0 {
    (ord_concat_ptr, self.ord_concat_dim)
} else {
    (h_s2_ptr, self.shared_h2)
};

For d == 3:

let (fc_input, fc_k) = if d == 3 && urg_concat_ptr != 0 {
    (urg_concat_ptr, self.urg_concat_dim)
} else {
    (h_s2_ptr, self.shared_h2)
};

• Step 7: Apply the same changes to forward_target_raw and forward_vsn_glu_branch.

• Step 8: Update the backward pass in batched_backward.rs to handle the wider input dimension for branches 2 and 3. The backward_fc_layer calls for w_b2fc and w_b3fc need in_dim = shared_h2 + 3 when concat is active. Also accumulate d_h_s2 from the first SH2 columns of the wider gradients (same accumulate_d_h_s2_from_concat pattern used for branch 1).

• Step 9: Build check:

SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -10

• Step 10: Commit:

git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs crates/ml/src/cuda_pipeline/batched_forward.rs crates/ml/src/cuda_pipeline/batched_backward.rs
git commit -m "feat: inject OFI microstructure features into order and urgency branches

Widen w_b2fc, w_b3fc, w_gate_2, w_gate_3 from [AH,SH2] to [AH,SH2+3].
Order branch gets [bid_ask_spread, depth_imbalance_L1, queue_pressure].
Urgency branch gets [spread_velocity, depth_change_rate, trade_arrival_rate].
3072 extra params. Same concat pattern as direction-conditioned magnitude.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>"

---

Task 7: xLSTM Temporal Context

Files:

• crates/ml/src/trainers/dqn/trainer/mod.rs
• crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
• crates/ml/src/trainers/dqn/fused_training.rs

CPU-only mLSTM cell that processes Q-stats every 50 training steps. Outputs an 8-dim context vector used by liquid tau, selectivity, and backtracking.

• Step 1: Add QLSTMCell struct to crates/ml/src/trainers/dqn/trainer/mod.rs, before the DQNTrainer struct:

/// Lightweight mLSTM cell for temporal Q-value context (xLSTM-inspired).
///
/// Processes 11-dim Q-stats input every 50 training steps.
/// Outputs an 8-dim context vector used by liquid tau, selectivity,
/// and trajectory backtracking.
///
/// All computation is CPU-side. Zero GPU overhead.
pub(crate) struct QLSTMCell {
    /// Weight matrices [11, 8] each
    w_k: [[f32; 8]; 11],
    w_v: [[f32; 8]; 11],
    w_q: [[f32; 8]; 11],
    w_i: [[f32; 8]; 11],
    w_f: [[f32; 8]; 11],
    w_o: [[f32; 8]; 11],
    /// Matrix memory [8, 8]
    c_matrix: [[f32; 8]; 8],
    /// Normalizer vector [8]
    n_vec: [f32; 8],
    /// Previous hidden state [8] (for divergence detection)
    h_prev: [f32; 8],
    /// Current context output [8]
    pub context: [f32; 8],
    /// Training: previous q_mean prediction for self-supervised loss
    pred_q_mean: f32,
    /// Learning rate for self-supervised weight update
    lr: f32,
}

impl QLSTMCell {
    pub(crate) fn new() -> Self {
        // Initialize weights with small random values (Xavier: sqrt(2/(11+8)) = 0.324)
        let mut cell = Self {
            w_k: [[0.0; 8]; 11],
            w_v: [[0.0; 8]; 11],
            w_q: [[0.0; 8]; 11],
            w_i: [[0.0; 8]; 11],
            w_f: [[0.0; 8]; 11],
            w_o: [[0.0; 8]; 11],
            c_matrix: [[0.0; 8]; 8],
            n_vec: [0.0; 8],
            h_prev: [0.0; 8],
            context: [0.0; 8],
            pred_q_mean: 0.0,
            lr: 0.001,
        };
        // Simple deterministic init: diagonal-dominant for stable initial behavior
        let scale = 0.324_f32;
        let mut seed = 0xCAFEBABEu64;
        let mut rand_f32 = |s: &mut u64| -> f32 {
            *s = s.wrapping_mul(6364136223846793005).wrapping_add(1);
            (*s >> 33) as f32 / (1u64 << 31) as f32 - 0.5
        };
        for w in [&mut cell.w_k, &mut cell.w_v, &mut cell.w_q,
                  &mut cell.w_i, &mut cell.w_f, &mut cell.w_o] {
            for i in 0..11 {
                for j in 0..8 {
                    w[i][j] = rand_f32(&mut seed) * scale;
                }
            }
        }
        // Forget gate bias: initialize w_f to produce ~0.9 sigmoid output (stable memory)
        for i in 0..11 {
            cell.w_f[i] = cell.w_f[i].map(|_| 0.1);
        }
        cell
    }

    /// Forward pass: process 11-dim Q-stats input, update matrix memory, produce 8-dim context.
    ///
    /// Input layout: [avg_max_q, q_min, q_max, q_mean, q_var, entropy, utilization, q_gap_0..3]
    pub(crate) fn forward(&mut self, q_input: &[f32; 11]) {
        // Save previous context for divergence detection
        self.h_prev = self.context;

        // Linear projections: k, v, q, i, f, o = W_x @ q_input
        let mut k_t = [0.0_f32; 8];
        let mut v_t = [0.0_f32; 8];
        let mut q_t = [0.0_f32; 8];
        let mut i_pre = [0.0_f32; 8];
        let mut f_pre = [0.0_f32; 8];
        let mut o_pre = [0.0_f32; 8];

        for j in 0..8 {
            for i in 0..11 {
                k_t[j] += self.w_k[i][j] * q_input[i];
                v_t[j] += self.w_v[i][j] * q_input[i];
                q_t[j] += self.w_q[i][j] * q_input[i];
                i_pre[j] += self.w_i[i][j] * q_input[i];
                f_pre[j] += self.w_f[i][j] * q_input[i];
                o_pre[j] += self.w_o[i][j] * q_input[i];
            }
        }

        // Gates
        let i_t: [f32; 8] = std::array::from_fn(|j| i_pre[j].exp().min(1e6)); // exponential input gate
        let f_t: [f32; 8] = std::array::from_fn(|j| 1.0 / (1.0 + (-f_pre[j]).exp())); // sigmoid forget
        let o_t: [f32; 8] = std::array::from_fn(|j| 1.0 / (1.0 + (-o_pre[j]).exp())); // sigmoid output

        // Matrix memory update: C_t = f_t * C_{t-1} + i_t * outer(v_t, k_t)
        for r in 0..8 {
            for c in 0..8 {
                self.c_matrix[r][c] = f_t[r] * self.c_matrix[r][c] + i_t[r] * v_t[r] * k_t[c];
            }
        }

        // Normalizer update: n_t = f_t * n_{t-1} + i_t * k_t
        for j in 0..8 {
            self.n_vec[j] = f_t[j] * self.n_vec[j] + i_t[j] * k_t[j];
        }

        // Retrieval: h_t = o_t * (C_t @ q_t) / max(|n_t @ q_t|, 1)
        let mut c_q = [0.0_f32; 8];
        for r in 0..8 {
            for c in 0..8 {
                c_q[r] += self.c_matrix[r][c] * q_t[c];
            }
        }

        let n_q: f32 = self.n_vec.iter().zip(q_t.iter()).map(|(n, q)| n * q).sum();
        let denom = n_q.abs().max(1.0);

        self.context = std::array::from_fn(|j| o_t[j] * c_q[j] / denom);
    }

    /// Context divergence: L2 distance between current and previous context.
    /// Used as additional plateau signal for trajectory backtracking.
    pub(crate) fn context_divergence(&self) -> f32 {
        self.context.iter().zip(self.h_prev.iter())
            .map(|(a, b)| (a - b) * (a - b))
            .sum::<f32>()
            .sqrt()
    }

    /// Context L2 norm. Used to trigger RK4 vs Euler switching.
    pub(crate) fn context_norm(&self) -> f32 {
        self.context.iter().map(|x| x * x).sum::<f32>().sqrt()
    }
}

• Step 2: Add qlstm_cell field to DQNTrainer struct (in mod.rs, after the saboteur field around line 347):

    /// xLSTM-inspired temporal Q-value context cell (CPU-only).
    pub(crate) qlstm_cell: QLSTMCell,

• Step 3: Initialize in the constructor (constructor.rs). Find where saboteur is initialized and add nearby:

            qlstm_cell: QLSTMCell::new(),

• Step 4: Wire the cell into the Q-stats processing path. In fused_training.rs, find where reduce_current_q_stats results are consumed (the training loop calls this and uses the returned QValueStatsResult). After Q-stats are downloaded to CPU, feed them to the QLSTM cell.

In the training step path (in training_loop.rs or wherever Q-stats are processed every 50 steps), add:

    // Feed Q-stats to xLSTM cell for temporal context
    let qlstm_input: [f32; 11] = [
        q_stats.avg_max_q,
        q_stats.q_min,
        q_stats.q_max,
        q_stats.q_mean,
        q_stats.q_var,
        q_stats.atom_entropy,
        q_stats.atom_utilization,
        q_stats.per_branch_q_gap[0],
        q_stats.per_branch_q_gap[1],
        q_stats.per_branch_q_gap[2],
        q_stats.per_branch_q_gap[3],
    ];
    self.qlstm_cell.forward(&qlstm_input);

• Step 5: Build check:

SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -10

• Step 6: Commit:

git add crates/ml/src/trainers/dqn/trainer/mod.rs crates/ml/src/trainers/dqn/trainer/constructor.rs crates/ml/src/trainers/dqn/trainer/training_loop.rs
git commit -m "feat: add xLSTM temporal Q-value context cell (CPU-only)

QLSTMCell with matrix memory processes 11-dim Q-stats every 50 steps.
Outputs 8-dim context vector. Provides context_divergence() for
backtracking and context_norm() for RK4/Euler switching.
528 CPU-side params (6 weight matrices [11,8]). Zero GPU overhead.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>"

---

Task 8: RK4 Adaptive ODE -- Replace Euler in update_liquid_tau

Files:

• crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs

Replace Euler integration in update_liquid_tau with adaptive RK4/Euler switching based on xLSTM context norm.

• Step 1: Add rk4_threshold field to GpuDqnTrainer (near the liquid tau fields around line 650):

    /// RK4/Euler switching threshold on xLSTM context norm.
    /// When context_norm > rk4_threshold, use RK4 for higher accuracy.
    rk4_threshold: f32,

Initialize to 0.5 in the constructor.

• Step 2: Add a set_qlstm_context_norm method for receiving the norm from the training loop:

    /// Set the current xLSTM context norm for RK4/Euler switching.
    pub fn set_qlstm_context_norm(&mut self, norm: f32) {
        self.qlstm_context_norm = norm;
    }

Add the field:

    qlstm_context_norm: f32,

Initialize to 0.0.

• Step 3: Replace the Euler step in update_liquid_tau (line 1101-1126) with adaptive RK4/Euler:

    pub fn update_liquid_tau(&mut self, per_branch_q_gaps: [f32; 4]) {
        let tau_min = 0.01_f32;
        let tau_max = 1.0_f32;
        let delta_z_approx = self.eval_q_std_ema.max(0.01);

        for d in 0..4 {
            let q_gap_d = per_branch_q_gaps[d];

            if self.per_branch_q_gap_ema[d] < 1e-12 {
                self.per_branch_q_gap_ema[d] = q_gap_d;
            } else {
                let err = (q_gap_d - self.per_branch_q_gap_ema[d]).abs();
                let alpha = (err / (err + self.per_branch_q_gap_ema[d].max(0.001))).clamp(0.01, 0.3);
                self.per_branch_q_gap_ema[d] = (1.0 - alpha) * self.per_branch_q_gap_ema[d] + alpha * q_gap_d;
            }

            let velocity_d = q_gap_d - self.per_branch_q_gap_ema[d];
            let sigmoid_arg = velocity_d / delta_z_approx;
            let sig = 1.0 / (1.0 + (-sigmoid_arg).exp());
            let tau_d = tau_min + (tau_max - tau_min) * sig;

            let current = unsafe { *self.liquid_mod_pinned.add(d) };

            // ODE: dx/dt = f(x) = (1/tau_d) * (1.0 - x)
            let f = |x: f32| -> f32 { (1.0 / tau_d) * (1.0 - x) };

            let updated = if self.qlstm_context_norm > self.rk4_threshold {
                // RK4 step (higher accuracy during regime transitions)
                let k1 = f(current);
                let k2 = f(current + 0.5 * k1);
                let k3 = f(current + 0.5 * k2);
                let k4 = f(current + k3);
                current + (k1 + 2.0 * k2 + 2.0 * k3 + k4) / 6.0
            } else {
                // Euler step (fast, accurate enough in stable regime)
                current + f(current)
            };

            unsafe { *self.liquid_mod_pinned.add(d) = updated.clamp(0.1, 2.0); }
        }
    }

• Step 4: Wire the context norm from the training loop. In fused_training.rs, add a passthrough:

    pub(crate) fn set_qlstm_context_norm(&mut self, norm: f32) {
        self.trainer.set_qlstm_context_norm(norm);
    }

In the training loop, after the QLSTM forward call, propagate the norm:

    if let Some(ref fused) = self.fused_training_ctx {
        fused.set_qlstm_context_norm(self.qlstm_cell.context_norm());
    }

• Step 5: Build check:

SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -10

• Step 6: Commit:

git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs crates/ml/src/trainers/dqn/fused_training.rs crates/ml/src/trainers/dqn/trainer/training_loop.rs
git commit -m "feat: replace Euler with adaptive RK4/Euler in liquid tau ODE

When xLSTM context_norm > 0.5 (regime transition detected), use 4th-order
Runge-Kutta for accurate integration. Otherwise keep Euler for speed.
Zero new params — pure algorithm change on existing CPU-side update path.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>"

---

Task 10: Q-Mean Drift Correction (Component 8) — EXECUTE FIRST

Priority: Execute BEFORE Tasks 1-9. Fixes active instability in live H100 run.

Files:

• crates/ml/src/cuda_pipeline/experience_kernels.cu

• crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs

• Step 1: Add q_mean_center kernel to experience_kernels.cu:

/**
 * Zero-mean Q-values: subtract global mean from all Q-values.
 * Prevents bootstrapping drift from accumulating across epochs.
 * Preserves action ranking (constant subtraction doesn't change argmax).
 *
 * Phase 1 (block 0): compute mean of Q_raw[B*total_actions]
 * Phase 2 (all blocks): Q_centered[i] = Q_raw[i] - mean
 *
 * Grid: ceil(B*total_actions/256), Block: 256.
 */
extern "C" __global__ void q_mean_center(
    float* __restrict__ q_values,        /* [B, total_actions] in-place */
    float* __restrict__ mean_scratch,    /* [1] scratch for global mean */
    int total                            /* B * total_actions */
) {
    /* Phase 1: block 0 computes mean */
    if (blockIdx.x == 0) {
        __shared__ float sdata[256];
        float sum = 0.0f;
        for (int i = threadIdx.x; i < total; i += blockDim.x) {
            sum += q_values[i];
        }
        sdata[threadIdx.x] = sum;
        __syncthreads();
        for (int s = blockDim.x / 2; s > 0; s >>= 1) {
            if (threadIdx.x < s) sdata[threadIdx.x] += sdata[threadIdx.x + s];
            __syncthreads();
        }
        if (threadIdx.x == 0) {
            mean_scratch[0] = sdata[0] / (float)total;
        }
    }
    /* Global barrier via grid sync not available — use two-phase launch */
}

Actually, this needs a two-phase approach (can't barrier across blocks). Use two kernel launches:

/**
 * Phase 1: Compute mean of Q-values into a scratch scalar.
 * Grid: (1, 1, 1), Block: (256, 1, 1).
 */
extern "C" __global__ void q_mean_reduce(
    const float* __restrict__ q_values,  /* [total] */
    float* __restrict__ mean_out,        /* [1] output */
    int total)
{
    __shared__ float sdata[256];
    float sum = 0.0f;
    for (int i = threadIdx.x; i < total; i += 256) {
        sum += q_values[i];
    }
    sdata[threadIdx.x] = sum;
    __syncthreads();
    for (int s = 128; s > 0; s >>= 1) {
        if (threadIdx.x < s) sdata[threadIdx.x] += sdata[threadIdx.x + s];
        __syncthreads();
    }
    if (threadIdx.x == 0) mean_out[0] = sdata[0] / (float)total;
}

/**
 * Phase 2: Subtract mean from all Q-values.
 * Grid: ceil(total/256), Block: 256.
 */
extern "C" __global__ void q_mean_subtract(
    float* __restrict__ q_values,        /* [total] in-place */
    const float* __restrict__ mean_val,  /* [1] */
    int total)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= total) return;
    q_values[i] -= mean_val[0];
}

• Step 2: Load kernels in gpu_dqn_trainer.rs constructor. Add fields q_mean_reduce_kernel, q_mean_subtract_kernel, q_mean_scratch: CudaSlice<f32> (1 float).

• Step 3: Add launch_q_mean_center method:

pub(crate) fn launch_q_mean_center(&self, batch_size: usize) -> Result<(), MLError> {
    let total = (batch_size * self.total_actions()) as i32;
    let q_ptr = self.q_out_buf.raw_ptr();
    let scratch_ptr = self.q_mean_scratch.raw_ptr();
    // Phase 1: reduce to mean
    unsafe {
        self.stream.launch_builder(&self.q_mean_reduce_kernel)
            .arg(&q_ptr).arg(&scratch_ptr).arg(&total)
            .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (256,1,1), shared_mem_bytes: 0 })
            .map_err(|e| MLError::ModelError(format!("q_mean_reduce: {e}")))?;
    }
    // Phase 2: subtract mean
    let blocks = ((total as u32 + 255) / 256).max(1);
    unsafe {
        self.stream.launch_builder(&self.q_mean_subtract_kernel)
            .arg(&q_ptr).arg(&scratch_ptr).arg(&total)
            .launch(LaunchConfig { grid_dim: (blocks,1,1), block_dim: (256,1,1), shared_mem_bytes: 0 })
            .map_err(|e| MLError::ModelError(format!("q_mean_subtract: {e}")))?;
    }
    Ok(())
}

• Step 4: Wire into Q-value pipeline. In reduce_current_q_stats, AFTER compute_expected_q and BEFORE launch_q_attention:

self.launch_q_mean_center(batch_size)?;

• Step 5: Add passthrough in fused_training.rs.

• Step 6: Verify build and commit:

SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5
git add crates/ml/src/cuda_pipeline/experience_kernels.cu \
       crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs \
       crates/ml/src/trainers/dqn/fused_training.rs
git commit -m "feat: Q-mean centering — prevents bootstrapping drift (Component 8)

Two-phase kernel: q_mean_reduce computes global mean, q_mean_subtract
centers all Q-values. Runs after compute_expected_q, before Q-attention.
Fixes Q-mean drift 0→+0.47 observed in train-vpb4w.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>"

---

Task 11: Atom Utilization Recovery (Component 9) — EXECUTE SECOND

Priority: Execute AFTER Task 10, BEFORE Tasks 1-9. Fixes atom collapse in live H100 run.

Files:

• crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs

• crates/ml/src/trainers/dqn/fused_training.rs

• Step 1: Add utilization_ema field to GpuDqnTrainer:

/// EMA of atom utilization [0,1] for adaptive entropy regularization.
utilization_ema: f32,

Initialize to 1.0 in the constructor (fully utilized at start).

• Step 2: In reduce_current_q_stats, after reading atom_utilization from the Q-stats readback, update the EMA:

// Adaptive-rate EMA for atom utilization
let util = stats.atom_utilization; // already in QValueStatsResult
if self.utilization_ema > 0.99 {
    self.utilization_ema = util;
} else {
    let err = (util - self.utilization_ema).abs();
    let alpha = (err / (err + 0.1)).clamp(0.01, 0.3);
    self.utilization_ema = (1.0 - alpha) * self.utilization_ema + alpha * util;
}

• Step 3: Compute adaptive entropy coefficient. In the method that launches c51_grad_kernel (or wherever entropy_coeff is prepared), replace the fixed hyperparameter:

// Adaptive entropy: ramp up when atom utilization drops below 50%
let base_entropy = self.config.entropy_coefficient;
let adaptive_entropy = if self.utilization_ema < 0.5 {
    base_entropy * (1.0 - self.utilization_ema / 0.5)  // linear ramp: 0 at 50%, base at 0%
} else {
    0.0  // no entropy bonus when atoms are healthy
};
let entropy_coeff = base_entropy + adaptive_entropy;

Find where entropy_coeff is passed to launch_c51_grad and replace the fixed value with the adaptive one.

• Step 4: Add accessor in fused_training.rs:

pub(crate) fn utilization_ema(&self) -> f32 { self.trainer.utilization_ema }

• Step 5: Verify build and commit:

SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5
git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs \
       crates/ml/src/trainers/dqn/fused_training.rs
git commit -m "feat: adaptive atom entropy — recovers utilization when C51 distribution collapses (Component 9)

entropy_coeff ramps up when atom utilization EMA drops below 50%.
Zero at 50%+, linear ramp to base_entropy at 0%. Prevents distributional
collapse observed in train-vpb4w (100%→20% over 18 epochs).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>"

---

Task 9: Build Verification + Smoke Tests

Files:

• All modified files from Tasks 1-8

Final verification that everything compiles and the smoke test passes.

• Step 1: Full workspace check:

SQLX_OFFLINE=true cargo check --workspace 2>&1 | tail -20

• Step 2: ML crate check:

SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -10

• Step 3: Run smoke tests (requires test data and CUDA GPU):

SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture 2>&1 | tail -40

• Step 4: If smoke tests fail, diagnose and fix. Common failure modes:

• NUM_WEIGHT_TENSORS mismatch: check all arrays sized to 50

• param_sizes sum overflow: verify total_params fits in GPU allocation

• Kernel launch failure: check argument count matches kernel signature

• NaN from spline eval: verify clamp bounds in kan_gate_combine

• Step 5: Verify parameter count increase is reasonable:

# Expected: ~15,120 new params
# KAN spline gates: 9,216 (4 branches x 256 x 9)
# Graph message passing: 60 (4 edges x 15)
# Diffusion refinement: 2,772 (3 steps x 924)
# TLOB widening: 3,072 (4 tensors x 768)
# xLSTM (CPU only): 528 (not in GPU param count)
# Quantile Q-select: 0
# RK4 ODE: 0
# Total GPU: 15,120 new params

• Step 6: Final commit with all verification:

git add -A
git status
# Verify no secrets or .env files are staged
git commit -m "chore: verify supervised architecture transfer builds and passes smoke tests

All 7 components integrated:
- Quantile Q-select (0 params, replaces Boltzmann)
- Graph message passing (60 params, 4 directed edges)
- KAN spline gates (9216 params, replaces GLU sigmoid)
- Diffusion Q-refinement (2772 params, 3-step denoiser)
- TLOB injection (3072 params, OFI→order/urgency branches)
- xLSTM context (528 CPU params, temporal Q-stats memory)
- RK4 adaptive ODE (0 params, replaces Euler)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>"