/**
 * Decision Transformer CUDA Kernels (Chen et al., 2021).
 *
 * Forward pass: embed → {QKV → causal attention → LN → FFN → LN} × N → action head → CE loss
 * Backward pass (Phase 1): CE backward → action head backward → residual passthrough → embed backward
 *
 * Compiled via NVRTC source concatenation (no #include).
 * The following constants are injected via NVRTC #define:
 *   DT_EMBED_DIM, DT_NUM_HEADS, DT_CONTEXT_LEN, DT_INPUT_DIM, DT_NUM_ACTIONS
 */

/* ── Compile-time defaults (overridden by NVRTC injection) ───────────── */
#ifndef DT_EMBED_DIM
#define DT_EMBED_DIM 128
#endif
#ifndef DT_NUM_HEADS
#define DT_NUM_HEADS 4
#endif
#ifndef DT_CONTEXT_LEN
#define DT_CONTEXT_LEN 20
#endif
#ifndef DT_INPUT_DIM
#define DT_INPUT_DIM 74
#endif
#ifndef DT_NUM_ACTIONS
#define DT_NUM_ACTIONS 9
#endif

#define DT_HEAD_DIM (DT_EMBED_DIM / DT_NUM_HEADS)
#define DT_FFN_DIM  (DT_EMBED_DIM * 4)

/* ══════════════════════════════════════════════════════════════════════
 * KERNEL 1: TOKEN EMBEDDING + POSITIONAL ENCODING
 *
 * Linear projection of each timestep's token (return-to-go, state, action)
 * plus learned positional encoding.
 *
 * out[b][t][d] = sum_j(traj[b][t][j] * W_embed[j][d]) + b_embed[d] + pos_embed[t][d]
 *
 * Grid:  (B*T, 1, 1)
 * Block: (min(E, 256), 1, 1)
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_embed_kernel(
    const float* __restrict__ trajectories,  /* [B, T, input_dim] */
    const float* __restrict__ W_embed,       /* [input_dim, embed_dim] */
    const float* __restrict__ b_embed,       /* [embed_dim] */
    const float* __restrict__ pos_embed,     /* [T, embed_dim] */
    float*       __restrict__ output,        /* [B, T, embed_dim] */
    int B, int T, int input_dim, int embed_dim
) {
    int bt = blockIdx.x;                /* combined (batch, timestep) index */
    if (bt >= B * T) return;
    int d = threadIdx.x;                /* embed dim index */

    int b = bt / T;
    int t = bt % T;

    const float* traj = trajectories + (b * T + t) * input_dim;
    float* out = output + (b * T + t) * embed_dim;

    for (int dd = d; dd < embed_dim; dd += blockDim.x) {
        float val = b_embed[dd] + pos_embed[t * embed_dim + dd];
        for (int j = 0; j < input_dim; j++) {
            val = val + traj[j] * W_embed[j * embed_dim + dd];
        }
        out[dd] = val;
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * KERNEL 2: QKV PROJECTION
 *
 * Three linear projections for multi-head attention.
 * Q[b][t][d] = sum_j(input[b][t][j] * W_Q[j][d])   (no bias for Q/K/V in standard DT)
 *
 * Grid:  (B*T, 1, 1)
 * Block: (min(E, 256), 1, 1)
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_qkv_projection_kernel(
    const float* __restrict__ input,    /* [B, T, E] */
    const float* __restrict__ W_Q,      /* [E, E] */
    const float* __restrict__ W_K,      /* [E, E] */
    const float* __restrict__ W_V,      /* [E, E] */
    float*       __restrict__ Q_out,    /* [B, T, E] */
    float*       __restrict__ K_out,    /* [B, T, E] */
    float*       __restrict__ V_out,    /* [B, T, E] */
    int B, int T, int E
) {
    int bt = blockIdx.x;
    if (bt >= B * T) return;
    int d = threadIdx.x;

    const float* x = input + bt * E;

    for (int dd = d; dd < E; dd += blockDim.x) {
        float q = 0.0f, k = 0.0f, v = 0.0f;
        for (int j = 0; j < E; j++) {
            float xj = x[j];
            q = q + xj * W_Q[j * E + dd];
            k = k + xj * W_K[j * E + dd];
            v = v + xj * W_V[j * E + dd];
        }
        Q_out[bt * E + dd] = q;
        K_out[bt * E + dd] = k;
        V_out[bt * E + dd] = v;
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * KERNEL 3: CAUSAL SELF-ATTENTION
 *
 * Per-head causal attention with softmax.
 * Grid:  (B, num_heads, 1) — one block per (batch, head)
 * Block: (T, 1, 1) — one thread per query timestep
 *
 * For each query position i:
 *   score[j] = Q[i] dot K[j] / sqrt(D_h)   for j <= i  (causal)
 *   score[j] = -inf                          for j > i
 *   attn[j] = softmax(score[j])
 *   out[i] = sum_j(attn[j] * V[j])
 *
 * Output projection W_O is also applied here:
 *   final[b][t][d] = sum_h( out_h[t][dh] * W_O[h*Dh + dh, d] )
 *
 * NOTE: Attention score computation uses F32 accumulation for numerical
 * stability (softmax is sensitive to precision). Results stored as BF16.
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_causal_attention_kernel(
    const float* __restrict__ Q,        /* [B, T, E] */
    const float* __restrict__ K,        /* [B, T, E] */
    const float* __restrict__ V,        /* [B, T, E] */
    const float* __restrict__ W_O,      /* [E, E] output projection */
    float*       __restrict__ per_head_out, /* [B, num_heads, T, E] per-head projected output */
    int B, int T, int E, int num_heads
) {
    int b = blockIdx.x;
    int h = blockIdx.y;
    if (b >= B || h >= num_heads) return;
    int t_idx = threadIdx.x;  /* query timestep */
    if (t_idx >= T) return;

    int Dh = E / num_heads;
    int head_off = h * Dh;

    /* Pointers for this batch element */
    const float* Q_b = Q + b * T * E;
    const float* K_b = K + b * T * E;
    const float* V_b = V + b * T * E;

    /* Shared memory for K and V of this head: [T, Dh] each */
    extern __shared__ float shmem[];
    float* sh_K = shmem;                /* [T, Dh] */
    float* sh_V = shmem + T * Dh;       /* [T, Dh] */

    /* Load K and V for this head into shared memory */
    for (int dd = 0; dd < Dh; dd++) {
        sh_K[t_idx * Dh + dd] = K_b[t_idx * E + head_off + dd];
        sh_V[t_idx * Dh + dd] = V_b[t_idx * E + head_off + dd];
    }
    __syncthreads();

    /* Compute attention scores Q[t_idx] dot K[j] / sqrt(Dh) for j <= t_idx
     * F32 accumulation for softmax numerical stability */
    float inv_sqrt_dh = rsqrtf((float)Dh);

    /* Numerically stable softmax: find max score first */
    float max_score = -1e30f;
    for (int j = 0; j <= t_idx; j++) {
        float score = 0.0f;
        for (int dd = 0; dd < Dh; dd++) {
            score += (float)Q_b[t_idx * E + head_off + dd] * (float)sh_K[j * Dh + dd];
        }
        score *= inv_sqrt_dh;
        if (score > max_score) max_score = score;
    }

    /* Compute softmax weights and weighted sum of V */
    float attn_out[128]; /* max Dh — static array to avoid dynamic alloc */
    for (int dd = 0; dd < Dh; dd++) attn_out[dd] = 0.0f;

    float sum_exp = 0.0f;
    for (int j = 0; j <= t_idx; j++) {
        float score = 0.0f;
        for (int dd = 0; dd < Dh; dd++) {
            score += (float)Q_b[t_idx * E + head_off + dd] * (float)sh_K[j * Dh + dd];
        }
        float weight = expf(score * inv_sqrt_dh - max_score);
        sum_exp += expf(score * inv_sqrt_dh - max_score);
        for (int dd = 0; dd < Dh; dd++) {
            attn_out[dd] = attn_out[dd] + weight * sh_V[j * Dh + dd];
        }
    }

    /* Normalize by softmax denominator */
    float inv_sum = 1.0f / (sum_exp + 1e-8f);
    for (int dd = 0; dd < Dh; dd++) {
        attn_out[dd] = attn_out[dd] * inv_sum;
    }

    /* Write per-head attention output to a per-head buffer.
     * A separate dt_sum_heads_kernel will combine all heads deterministically. */

    /* per_head_out layout: [B, num_heads, T, E] — each head writes its
     * projected output to its own slice, avoiding cross-head atomicAdd. */
    float* head_slice = per_head_out + ((b * num_heads + h) * T + t_idx) * E;

    /* Output projection: head_out[dh] -> projected[d] via W_O
     * out[d] = sum_dh(attn_out[dh] * W_O[(h*Dh+dh)*E + d]) */
    for (int dd = 0; dd < E; dd++) {
        float proj = 0.0f;
        for (int dh = 0; dh < Dh; dh++) {
            proj = proj + attn_out[dh] * W_O[(head_off + dh) * E + dd];
        }
        head_slice[dd] = proj;
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * KERNEL 3b: SUM PER-HEAD ATTENTION OUTPUTS
 *
 * Sums the per-head projected outputs into the final attention output
 * and adds the output bias. Fully deterministic — fixed summation order.
 *
 * output[b][t][d] = b_O[d] + sum_h(per_head_out[b][h][t][d])
 *
 * Grid:  (B*T, 1, 1)
 * Block: (min(E, 256), 1, 1)
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_sum_heads_kernel(
    const float* __restrict__ per_head_out,  /* [B, num_heads, T, E] */
    const float* __restrict__ b_O,           /* [E] output bias */
    float*       __restrict__ output,        /* [B, T, E] */
    int B, int T, int E, int num_heads
) {
    int bt = blockIdx.x;
    if (bt >= B * T) return;
    int d = threadIdx.x;

    int b = bt / T;
    int t = bt % T;

    for (int dd = d; dd < E; dd += blockDim.x) {
        float sum = b_O[dd];
        for (int h = 0; h < num_heads; h++) {
            sum += per_head_out[((b * num_heads + h) * T + t) * E + dd];
        }
        output[bt * E + dd] = sum;
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * KERNEL 4: LAYER NORMALIZATION
 *
 * output = gamma * (input - mean) / sqrt(var + eps) + beta
 *
 * NOTE: Mean/variance computation uses F32 accumulation for numerical
 * stability (reduction over large vectors). Final output stored as BF16.
 *
 * Grid:  (B*T, 1, 1)
 * Block: (min(E, 256), 1, 1)
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_layernorm_kernel(
    const float* __restrict__ input,    /* [B*T, E] */
    const float* __restrict__ gamma,    /* [E] */
    const float* __restrict__ beta,     /* [E] */
    float*       __restrict__ output,   /* [B*T, E] */
    int N,  /* B*T */
    int E
) {
    int n = blockIdx.x;
    if (n >= N) return;
    int tid = threadIdx.x;

    const float* x = input + n * E;
    float* out = output + n * E;

    /* Compute mean via block-parallel reduction (F32 for precision) */
    float local_sum = 0.0f;
    for (int d = tid; d < E; d += blockDim.x) {
        local_sum += (float)x[d];
    }

    /* Warp-level reduction */
    for (int offset = 16; offset > 0; offset >>= 1)
        local_sum += __shfl_xor_sync(0xFFFFFFFF, local_sum, offset);

    /* Cross-warp reduction via shared memory */
    __shared__ float warp_sums[8]; /* up to 8 warps */
    int warp_id = tid / 32;
    int lane = tid % 32;
    if (lane == 0 && warp_id < 8) warp_sums[warp_id] = local_sum;
    __syncthreads();

    float total_sum;
    if (warp_id == 0) {
        float val = (lane < (int)((blockDim.x + 31) / 32)) ? warp_sums[lane] : 0.0f;
        for (int offset = 16; offset > 0; offset >>= 1)
            val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
        total_sum = val;
    }
    total_sum = __shfl_sync(0xFFFFFFFF, total_sum, 0);
    /* Broadcast from warp 0 lane 0 to all threads */
    __shared__ float sh_mean;
    if (tid == 0) sh_mean = total_sum / (float)E;
    __syncthreads();
    float mean = sh_mean;

    /* Compute variance */
    float local_var = 0.0f;
    for (int d = tid; d < E; d += blockDim.x) {
        float diff = (float)x[d] - mean;
        local_var += diff * diff;
    }

    for (int offset = 16; offset > 0; offset >>= 1)
        local_var += __shfl_xor_sync(0xFFFFFFFF, local_var, offset);

    if (lane == 0 && warp_id < 8) warp_sums[warp_id] = local_var;
    __syncthreads();

    float total_var;
    if (warp_id == 0) {
        float val = (lane < (int)((blockDim.x + 31) / 32)) ? warp_sums[lane] : 0.0f;
        for (int offset = 16; offset > 0; offset >>= 1)
            val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
        total_var = val;
    }
    total_var = __shfl_sync(0xFFFFFFFF, total_var, 0);
    __shared__ float sh_inv_std;
    if (tid == 0) sh_inv_std = rsqrtf(total_var / (float)E + 1e-5f);
    __syncthreads();
    float inv_std = sh_inv_std;

    /* Normalize + affine — final store as BF16 */
    for (int d = tid; d < E; d += blockDim.x) {
        out[d] = (float)gamma[d] * ((float)x[d] - mean) * inv_std + (float)beta[d];
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * KERNEL 5: FFN (FEED-FORWARD NETWORK)
 *
 * hidden = GELU(input × W1 + b1)
 * output = hidden × W2 + b2 + input (residual)
 *
 * GELU approximation: 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
 *
 * NOTE: GELU computation uses F32 for the intermediate hidden activations
 * (tanh precision matters). Final output stored as BF16.
 *
 * Grid:  (B*T, 1, 1)
 * Block: (min(E, 256), 1, 1)
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_ffn_kernel(
    const float* __restrict__ input,    /* [B*T, E] */
    const float* __restrict__ W1,       /* [E, 4E] */
    const float* __restrict__ b1,       /* [4E] */
    const float* __restrict__ W2,       /* [4E, E] */
    const float* __restrict__ b2,       /* [E] */
    float*       __restrict__ output,   /* [B*T, E] */
    int N,  /* B*T */
    int E
) {
    int n = blockIdx.x;
    if (n >= N) return;
    int tid = threadIdx.x;

    int FFN = E * 4;
    const float* x = input + n * E;
    float* out = output + n * E;

    /* Each thread computes its output elements via stride loop.
     * For the intermediate FFN dimension (4E), we need to compute
     * the full hidden vector. Use shared memory for the hidden activations.
     * F32 accumulation for GELU precision. */
    extern __shared__ float sh_hidden[];  /* [4E] — float accumulator for GELU */

    /* Step 1: Compute hidden = GELU(x @ W1 + b1) */
    for (int ff = tid; ff < FFN; ff += blockDim.x) {
        float h = (float)b1[ff];
        for (int j = 0; j < E; j++) {
            h += (float)x[j] * (float)W1[j * FFN + ff];
        }
        /* GELU approximation */
        float x3 = h * h * h;
        float inner = 0.7978845608f * (h + 0.044715f * x3);  /* sqrt(2/pi) ≈ 0.7979 */
        h = 0.5f * h * (1.0f + tanhf(inner));
        sh_hidden[ff] = h;
    }
    __syncthreads();

    /* Step 2: output = hidden @ W2 + b2 + input (residual) */
    for (int d = tid; d < E; d += blockDim.x) {
        float val = (float)b2[d] + (float)x[d];  /* bias + residual */
        for (int ff = 0; ff < FFN; ff++) {
            val += sh_hidden[ff] * (float)W2[ff * E + d];
        }
        out[d] = val;
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * KERNEL 6: ACTION HEAD (LINEAR PROJECTION)
 *
 * logits[b][t][a] = sum_d(input[b][t][d] * W_head[d][A]) + b_head[a]
 *
 * Takes the last timestep's embedding and projects to action logits.
 *
 * Grid:  (B*T, 1, 1)
 * Block: (min(A, 256), 1, 1)
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_action_head_kernel(
    const float* __restrict__ input,    /* [B*T, E] */
    const float* __restrict__ W_head,   /* [E, A] */
    const float* __restrict__ b_head,   /* [A] */
    float*       __restrict__ logits,   /* [B*T, A] */
    int N,  /* B*T */
    int E,
    int A   /* num_actions */
) {
    int n = blockIdx.x;
    if (n >= N) return;
    int tid = threadIdx.x;

    const float* x = input + n * E;
    float* out = logits + n * A;

    for (int a = tid; a < A; a += blockDim.x) {
        float val = b_head[a];
        for (int d = 0; d < E; d++) {
            val = val + x[d] * W_head[d * A + a];
        }
        out[a] = val;
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * KERNEL 7: CROSS-ENTROPY LOSS
 *
 * Numerically stable cross-entropy:
 *   loss = -logits[target] + log(sum(exp(logits - max)))  + max
 *
 * NOTE: Loss computation uses F32 for numerical stability (log-sum-exp).
 * Per-sample loss and total loss stored as BF16.
 *
 * Grid:  (B*T, 1, 1)
 * Block: (min(A, 256), 1, 1) or (32, 1, 1) for warp reduction
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_cross_entropy_kernel(
    const float* __restrict__ logits,       /* [B*T, A] */
    const int*   __restrict__ targets,      /* [B*T] */
    float* __restrict__ per_sample_loss, /* [B*T] */
    int N,  /* B*T */
    int A   /* num_actions */
) {
    int n = blockIdx.x;
    if (n >= N) return;

    const float* lg = logits + n * A;
    int target = targets[n];

    /* Find max for numerical stability (F32) */
    float max_val = -1e30f;
    for (int a = 0; a < A; a++) {
        float lga = (float)lg[a];
        if (lga > max_val) max_val = lga;
    }

    /* log-sum-exp (F32) */
    float sum_exp = 0.0f;
    for (int a = 0; a < A; a++) {
        sum_exp += expf((float)lg[a] - max_val);
    }

    float log_sum_exp = max_val + logf(sum_exp + 1e-8f);
    float target_logit = (target >= 0 && target < A) ? (float)lg[target] : 0.0f;
    float loss = log_sum_exp - target_logit;

    /* Clamp to prevent NaN propagation */
    loss = fminf(fmaxf(loss, 0.0f), 100.0f);

    per_sample_loss[n] = loss;
    /* Total loss is computed by a separate dt_reduce_loss_kernel for
     * full determinism — no atomicAdd across samples. */
}

/* ══════════════════════════════════════════════════════════════════════
 * KERNEL 7b: REDUCE PER-SAMPLE LOSS TO TOTAL LOSS
 *
 * Deterministic sequential sum of per-sample losses to compute mean.
 * Single-thread kernel — N is small (B*T, typically 256*20=5120) so
 * the serial sum is fast and perfectly reproducible.
 *
 * Grid:  (1, 1, 1)
 * Block: (1, 1, 1)
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_reduce_loss_kernel(
    const float* __restrict__ per_sample_loss, /* [N] */
    float*       __restrict__ total_loss,      /* [1] */
    int N
) {
    float sum = 0.0f;
    for (int i = 0; i < N; i++) {
        sum += per_sample_loss[i];
    }
    total_loss[0] = sum / (float)N;
}

/* ══════════════════════════════════════════════════════════════════════
 * KERNEL 8: CROSS-ENTROPY BACKWARD
 *
 * dL/d_logits = softmax(logits) - one_hot(target)
 *
 * NOTE: Softmax computation uses F32 for numerical stability.
 * Gradient output stored as BF16.
 *
 * Grid:  (B*T, 1, 1)
 * Block: (min(A, 256), 1, 1)
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_ce_backward_kernel(
    const float* __restrict__ logits,       /* [B*T, A] */
    const int*   __restrict__ targets,      /* [B*T] */
    float* __restrict__ d_logits,     /* [B*T, A] */
    int N,  /* B*T */
    int A   /* num_actions */
) {
    int n = blockIdx.x;
    if (n >= N) return;
    int tid = threadIdx.x;

    const float* lg = logits + n * A;
    float* d_lg = d_logits + n * A;
    int target = targets[n];

    /* Numerically stable softmax (F32) */
    float max_val = -1e30f;
    for (int a = 0; a < A; a++) {
        float lga = (float)lg[a];
        if (lga > max_val) max_val = lga;
    }

    float sum_exp = 0.0f;
    for (int a = 0; a < A; a++) {
        sum_exp += expf((float)lg[a] - max_val);
    }
    float inv_sum = 1.0f / (sum_exp + 1e-8f);

    /* d_logits = softmax - one_hot, scaled by 1/N for mean loss */
    float scale = 1.0f / (float)N;
    for (int a = tid; a < A; a += blockDim.x) {
        float softmax_a = expf((float)lg[a] - max_val) * inv_sum;
        float one_hot = (a == target) ? 1.0f : 0.0f;
        d_lg[a] = (softmax_a - one_hot) * scale;
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * KERNEL 9: LINEAR BACKWARD — INPUT GRADIENT
 *
 * Computes the input gradient for a linear layer:
 *   d_input[n][i] = sum_o(d_output[n][o] * W[i][o])
 *
 * This is per-sample with no cross-sample accumulation — fully
 * deterministic by construction.
 *
 * Grid:  (N, 1, 1)
 * Block: (min(I, 256), 1, 1)
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_linear_backward_kernel(
    const float* __restrict__ d_output, /* [N, O] */
    const float* __restrict__ input,    /* [N, I] (unused — kept for API compat) */
    const float* __restrict__ W,        /* [I, O] */
    float* __restrict__ dW,       /* [I, O] (unused — computed by dt_linear_grad_kernel) */
    float* __restrict__ db,       /* [O]    (unused — computed by dt_linear_grad_kernel) */
    float* __restrict__ d_input,  /* [N, I] */
    int N,
    int I,
    int O
) {
    int n = blockIdx.x;
    if (n >= N) return;
    int tid = threadIdx.x;

    const float* dout = d_output + n * O;
    float* dx = d_input + n * I;

    /* d_input[n][i] = sum_o(d_output[n][o] * W[i][o]) */
    for (int i = tid; i < I; i += blockDim.x) {
        float val = 0.0f;
        for (int o = 0; o < O; o++) {
            val = val + dout[o] * W[i * O + o];
        }
        dx[i] = val;
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * KERNEL 9b: LINEAR BACKWARD — WEIGHT + BIAS GRADIENTS (DETERMINISTIC)
 *
 * One thread per weight parameter. Each thread loops over all N samples
 * and accumulates the gradient contribution in a register — fully
 * deterministic with zero atomics.
 *
 *   dW[i][o] = sum_n(input[n][i] * d_output[n][o])
 *   db[o]    = sum_n(d_output[n][o])
 *
 * Weight grid:  (ceil(I*O / 256), 1, 1)
 * Weight block: (256, 1, 1)
 *
 * Bias is computed by threads in the first ceil(O/256) blocks where
 * the thread's linear index maps to a valid bias element (idx < O).
 * Since I*O >> O, we piggyback the bias computation onto the weight
 * kernel: thread idx maps to a (i, o) pair, and if i==0 we also
 * compute db[o]. This avoids a separate kernel launch.
 *
 * For DT dimensions (E=128, A=9 -> I*O = 1152 or E=128, input_dim=74
 * -> I*O = 9472), this fits in a handful of blocks. The loop over N
 * (N = B*T, typically 5120) is the inner loop — cache-friendly on
 * d_output (sequential reads of O floats per sample).
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_linear_grad_kernel(
    const float* __restrict__ d_output, /* [N, O] */
    const float* __restrict__ input,    /* [N, I] */
    float*       __restrict__ dW,       /* [I, O] */
    float*       __restrict__ db,       /* [O] */
    int N,
    int I,
    int O
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total_weights = I * O;

    /* ── Weight gradient: one thread per (i, o) pair ── */
    if (idx < total_weights) {
        int i = idx / O;
        int o = idx % O;

        float sum_w = 0.0f;
        for (int n = 0; n < N; n++) {
            sum_w += input[n * I + i] * d_output[n * O + o];
        }
        dW[idx] = sum_w;

        /* ── Bias gradient: piggyback on i==0 threads ──
         * When i==0, this thread also computes db[o] = sum_n(d_output[n][o]).
         * Since O is small, all bias elements are covered by the first O threads. */
        if (i == 0) {
            float sum_b = 0.0f;
            for (int n = 0; n < N; n++) {
                sum_b += d_output[n * O + o];
            }
            db[o] = sum_b;
        }
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * KERNEL 10: GRADIENT ZERO
 *
 * Zeroes a buffer before backward accumulation.
 *
 * Grid:  (ceil(n/256), 1, 1)
 * Block: (256, 1, 1)
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_zero_kernel(
    float* __restrict__ buf,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) buf[i] = 0.0f;
}

/* ══════════════════════════════════════════════════════════════════════
 * KERNEL 11: RESIDUAL ADD
 *
 * output[i] += residual[i]
 *
 * Used after LayerNorm to add the residual connection.
 *
 * Grid:  (ceil(n/256), 1, 1)
 * Block: (256, 1, 1)
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_residual_add_kernel(
    float* __restrict__ output,
    const float* __restrict__ residual,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) output[i] = output[i] + residual[i];
}

/* ══════════════════════════════════════════════════════════════════════
 * KERNEL 12: RETURN-TO-GO (REVERSE CUMULATIVE SUM)
 *
 * R_t = r_t + gamma * R_{t+1}
 * Computed backward from T-1 to 0.  Each thread handles one episode.
 *
 * Grid:  (num_episodes, 1, 1)
 * Block: (1, 1, 1)
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_return_to_go_kernel(
    const float* __restrict__ rewards,       /* [num_episodes, episode_len] */
    float*       __restrict__ returns_to_go, /* [num_episodes, episode_len] */
    float gamma,
    int episode_len,
    int num_episodes
) {
    int ep = blockIdx.x;
    if (ep >= num_episodes) return;

    const float* r = rewards + ep * episode_len;
    float*       rtg = returns_to_go + ep * episode_len;

    /* Backward scan: R_{T-1} = r_{T-1}, R_t = r_t + gamma * R_{t+1} */
    float cumul = 0.0f;
    for (int t = episode_len - 1; t >= 0; t--) {
        cumul = r[t] + gamma * cumul;
        rtg[t] = cumul;
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * KERNEL 13: BUILD TRAJECTORY TOKENS
 *
 * Packs [return-to-go, state_features..., action] into the trajectory
 * tensor expected by dt_embed_kernel.
 *
 * For each (episode, timestep):
 *   out[ep * T * input_dim + t * input_dim + 0]              = rtg[ep * T + t]
 *   out[ep * T * input_dim + t * input_dim + 1..state_dim+1] = features[bar_idx * feat_dim + 0..feat_dim]
 *   out[ep * T * input_dim + t * input_dim + state_dim+1]    = action (as float)
 *
 * bar_idx = episode_start_indices[ep] + t
 *
 * Grid:  (num_episodes, 1, 1)
 * Block: (min(T, 256), 1, 1) — one thread per timestep within episode
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_build_trajectories_kernel(
    const float* __restrict__ features,              /* [num_bars, feat_dim] */
    const float* __restrict__ returns_to_go,         /* [num_episodes, T] */
    const int*   __restrict__ actions,               /* [num_bars] — global per-bar actions */
    const int*   __restrict__ episode_start_indices,  /* [num_episodes] */
    float*       __restrict__ trajectories,          /* [num_episodes, T, input_dim] */
    int*         __restrict__ target_actions,        /* [num_episodes * T] */
    int T,
    int feat_dim,
    int input_dim,
    int num_episodes,
    int num_bars
) {
    int ep = blockIdx.x;
    if (ep >= num_episodes) return;

    int start = episode_start_indices[ep];

    for (int t = threadIdx.x; t < T; t += blockDim.x) {
        int bar_idx = start + t;
        /* Clamp to valid bar range */
        if (bar_idx >= num_bars) bar_idx = num_bars - 1;

        float* out = trajectories + (ep * T + t) * input_dim;

        /* Slot 0: return-to-go */
        out[0] = returns_to_go[ep * T + t];

        /* Slots 1..feat_dim: state features */
        const float* feat = features + bar_idx * feat_dim;
        for (int f = 0; f < feat_dim; f++) {
            out[1 + f] = feat[f];
        }

        /* Slot state_dim+1: action as float (state_dim = feat_dim) */
        int act = actions[bar_idx];  /* index into global bar actions */
        out[feat_dim + 1] = (float)act;

        /* Target action for cross-entropy loss */
        target_actions[ep * T + t] = act;
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * KERNEL 14: COMPUTE BAR REWARDS + EXPERT ACTIONS
 *
 * Per bar: reward = close[t+1] / close[t] - 1
 * Expert action: if reward > threshold → Long(8), < -threshold → Short(0), else Flat(4)
 *
 * NOTE: Price ratio computation uses F32 for division precision.
 *
 * Grid:  (ceil(num_bars / 256), 1, 1)
 * Block: (256, 1, 1)
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_compute_rewards_actions_kernel(
    const float* __restrict__ targets,    /* [num_bars, TARGET_DIM=6] — fxcache targets */
    float*       __restrict__ rewards,    /* [num_bars] */
    int*         __restrict__ actions,    /* [num_bars] */
    int num_bars,
    float action_threshold
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= num_bars) return;

    /* Targets stride per `crates/ml/src/fxcache.rs`
     * (TARGET_DIM=6): col 0=preproc_close, 1=preproc_next, 2=raw_close,
     * 3=raw_next (FUTURE), 4=raw_open, 5=mid_open. Read raw_close at col 2.
     *
     * Pre-`063fd2716` the stride was 4 and this kernel read `i * 4 + 3`
     * (raw_next column) — the stale "OHLC" docstring was incorrect: the
     * fxcache never had an OHLC column layout, only the close-row variants.
     * Reading at the old stride against the new stride-6 buffer slid every
     * lookup into the wrong bar's data. */
    float close_t  = (float)targets[i * 6 + 2];

    float reward;
    if (i < num_bars - 1) {
        float close_t1 = (float)targets[(i + 1) * 6 + 2];
        /* Avoid division by zero */
        reward = (close_t > 1e-8f) ? (close_t1 / close_t - 1.0f) : 0.0f;
    } else {
        reward = 0.0f;  /* Last bar has no future */
    }
    rewards[i] = reward;

    /* Expert action: Long=8 (max exposure), Short=0 (min exposure), Flat=4 (neutral) */
    int action;
    if (reward > action_threshold) {
        action = 8;  /* Long — b0_size-1 */
    } else if (reward < -action_threshold) {
        action = 0;  /* Short */
    } else {
        action = 4;  /* Flat — b0_size/2 */
    }
    actions[i] = action;
}
