kernels.cu

#ifndef PUFFERLIB_KERNELS_CU
#define PUFFERLIB_KERNELS_CU

/* Kernels must launch on the current torch stream to be traced by cudagraphs.
 * Launch functions take cudaStream_t as parameter - callers (modules.cu) should
 * pass at::cuda::getCurrentCUDAStream() when using with torch.
 */

#include <cuda_runtime.h>
#include "ops.cuh"
#include <cuda_fp16.h>
#include <cuda_bf16.h>
#include <curand_kernel.h>
#include <c10/util/BFloat16.h>

#include <cstdio>
#include <cstdint>

#define SEQ_SIZE 256
#define BLOCK_SIZE 256
inline int grid_size(int N) {
    return (N + BLOCK_SIZE - 1) / BLOCK_SIZE;
}
inline int seq_size(int N) {
    return (N + SEQ_SIZE - 1) / SEQ_SIZE;
}

// If you can get this to work, go ahead. I tried.
// NVCC won't parse templated types in kernel launches
/*
template <template <class> class KernelFn, typename... Args>
void dispatch_and_launch(const at::Tensor& example_tensor, Args... args) {
    const int64_t N = example_tensor.numel();
    const int64_t block = LAUNCH_BLOCK_SIZE;
    const int64_t grid = (N + block - 1) / block;
    auto stream = at::cuda::getCurrentCUDAStream();
    at::cuda::CUDAGuard device_guard(example_tensor.device());

    at::ScalarType dtype = example_tensor.scalar_type();
    if (dtype == at::ScalarType::Float) {
        KernelFn<float><<<grid, block, 0, stream>>>(args..., N);
    } else if (dtype == at::ScalarType::Half) {
        KernelFn<__half><<<grid, block, 0, stream>>>(args..., N);
    } else if (dtype == at::ScalarType::BFloat16) {
        KernelFn<__nv_bfloat16><<<grid, block, 0, stream>>>(args..., N);
    } else {
        AT_ERROR("Unsupported dtype: ", dtype);
    }
}
*/

template<typename T>
__global__ void rmsnorm_forward_kernel(
    T* __restrict__ out,
    float* __restrict__ inv_norm_buf,
    const T* __restrict__ x,
    const T* __restrict__ weight,
    double eps,
    int T_total,
    int H,
    int B
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= B * T_total) return;

    int b = idx / T_total;
    int t = idx % T_total;
    int base = b*T_total*H + t*H;

    float sum_sq = 0.0f;
    for (int h = 0; h < H; h++) {
        int curr = base + h;
        float x_val = float(x[curr]);
        sum_sq += x_val * x_val;
    }

    float rms = sqrtf(sum_sq/H + eps);
    float inv_rms = 1.0f / rms;
    inv_norm_buf[idx] = inv_rms;

    for (int h = 0; h < H; h++) {
        int curr = base + h;
        out[curr] = T(weight[h] * x[curr] * inv_rms);
    }
}

template<typename T>
__global__ void rmsnorm_backward_kernel(
    T* __restrict__ grad_x,
    T* __restrict__ grad_weight,
    const T* __restrict__ grad_out,
    const float* __restrict__ inv_norm_buf,
    const T* __restrict__ x_buf,
    const T* __restrict__ weight,
    double eps,
    int T_total,
    int H,
    int B
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= T_total*H*B) return;
    int base = idx % H;
    int norm_idx = idx / H;

    float inv_rms = inv_norm_buf[norm_idx];
    float inv_rms_3 = inv_rms * inv_rms * inv_rms;

    grad_x[idx] = weight[base] * grad_out[idx] * inv_rms;
    grad_weight[idx] = grad_out[idx] * inv_rms;

    float wg_x = 0.0f;
    for (int h=0; h<H; h++) {
        float x = x_buf[base + h];
        float w = weight[h];
        float g = grad_out[base + h];
        wg_x += w*g*x;
    }
    float x = x_buf[idx];
    grad_x[idx] -= x*wg_x*inv_rms_3/float(H);
}

/*
template<typename T>
__global__ void rmsnorm_backward_kernel(
    T* grad_x,
    T* grad_weight,
    const T* grad_out,
    const float* inv_norm_buf,
    const T* x,
    const T* weight,
    double eps,
    int T_total,
    int H,
    int B
) {
    int total_elements = B * T_total * H;
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total_elements) return;

    int h = idx % H;
    int vec_idx = idx / H;                    // index of the vector (b,t)
    int offset = vec_idx * H;

    float inv_rms = inv_norm_buf[vec_idx];
    float inv_rms3 = inv_rms * inv_rms * inv_rms;

    // ∂L/∂γ_h += grad_out * (x / rms)
    float gw = grad_out[idx] * (float)x[idx] * inv_rms;
    atomicAdd((float*)&grad_weight[h], gw);

    // Compute reduction: sum_h weight[h] * grad_out[h] * x[h]
    float sum = 0.0f;
    for (int i = 0; i < H; ++i) {
        sum += (float)weight[i] * (float)grad_out[offset + i] * (float)x[offset + i];
    }
    float reduction = sum * inv_rms;  // = σ γ g hat_x

    float dx = (float)weight[h] * (float)grad_out[idx] * inv_rms
               - (float)x[idx] * reduction * inv_rms3 / H;

    grad_x[idx] = T(dx);
}
*/

template<typename T>
void launch_rmsnorm_forward(
    T* __restrict__ out,
    float* __restrict__ inv_norm_buf,
    const T* __restrict__ x,
    const T* __restrict__ weight,
    double eps,
    int T_total,
    int H,
    int B,
    cudaStream_t stream
) {
    int total = B * T_total;
    int grid = grid_size(total);

    rmsnorm_forward_kernel<T><<<grid, BLOCK_SIZE, 0, stream>>>(
        out,
        inv_norm_buf,
        x,
        weight,
        eps,
        T_total,
        H,
        B
    );

    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        fprintf(stderr, "CUDA kernel launch error in forward: %s\n", cudaGetErrorString(err));
    }
}

template<typename T>
void launch_rmsnorm_backward(
    T* __restrict__ grad_x,
    T* __restrict__ grad_weight,
    const T* __restrict__ grad_out,
    const float* __restrict__ inv_norm_buf,
    const T* __restrict__ x_buf,
    const T* __restrict__ weight,
    double eps,
    int T_total,
    int H,
    int B,
    cudaStream_t stream
) {
    // The backward is fully parallel
    // since the inv norm is cached
    int total = B * T_total * H;
    int grid = grid_size(total);

    rmsnorm_backward_kernel<T><<<grid, BLOCK_SIZE, 0, stream>>>(
        grad_x,
        grad_weight,
        grad_out,
        inv_norm_buf,
        x_buf,
        weight,
        eps,
        T_total,
        H,
        B
    );

    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        fprintf(stderr, "CUDA kernel launch error in backward: %s\n", cudaGetErrorString(err));
    }
}


// Fused kernel: chunk + mingru_gate + sigmoid(proj) * out
// combined is (B, 1, 3*H) containing [hidden, gate, proj] concatenated on last dim
// state is (B, 1, H)
// out is (B, 1, H) = sigmoid(proj) * mingru_out (final output)
// next_state is (B, 1, H) = mingru_out (recurrent state, without proj)
template<typename T>
__global__ void mingru_gate_inference_kernel(
    T* out,
    T* next_state,
    const T* combined,    // (B, 1, 3*H) = [hidden, gate, proj]
    const T* state_in,    // (B, 1, H)
    int H,
    int B
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int N = B * H;
    if (idx >= N) return;

    int b = idx / H;
    int h = idx % H;

    // Read from combined: layout is [hidden(H), gate(H), proj(H)] for each batch
    int combined_base = b * 3 * H;
    float hidden = float(combined[combined_base + h]);
    float gate = float(combined[combined_base + H + h]);
    float proj = float(combined[combined_base + 2 * H + h]);

    float state = float(state_in[idx]);

    // mingru_gate computation
    float gate_sigmoid = sigmoid(gate);
    float hidden_tilde = tilde_relu_fwd(hidden);
    float mingru_out = lerp(state, hidden_tilde, gate_sigmoid);

    // next_state is mingru_out (for recurrence)
    next_state[idx] = T(mingru_out);

    // out is sigmoid(proj) * mingru_out (final output)
    float proj_sigmoid = sigmoid(proj);
    out[idx] = T(proj_sigmoid * mingru_out);
}

template<typename T>
void launch_mingru_gate_inference(
    T* out,
    T* next_state,
    const T* combined,
    const T* state_in,
    int H,
    int B,
    cudaStream_t stream
) {
    int N = B * H;
    int grid = grid_size(N);
    mingru_gate_inference_kernel<T><<<grid, BLOCK_SIZE, 0, stream>>>(
        out,
        next_state,
        combined,
        state_in,
        H,
        B
    );

    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        fprintf(stderr, "CUDA kernel launch error: %s\n", cudaGetErrorString(err));
    }
}


template<typename T>
__global__ void log_coeffs_and_values_kernel(
    T* log_coeffs,
    T* log_values,
    const T* gate,
    const T* hidden,
    int N
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= N) return;

    float g = float(gate[idx]);
    float h = float(hidden[idx]);

    log_coeffs[idx] = -softplus_fwd(g);
    float log_z = -softplus_fwd(-g);
    float log_tilde_h;
    if (h >= 0.0f) {
        float relu_h = relu(h);
        log_tilde_h = logf(relu_h + 0.5f);
    } else {
        log_tilde_h = -softplus_fwd(-h);
    }
    log_values[idx] = log_z + log_tilde_h;
}

template<typename T>
__global__ void log_coeffs_and_values_backward_kernel(
    T* grad_gate,
    T* grad_hidden,
    const T* grad_log_coeffs,
    const T* grad_log_values,
    const T* gate,
    const T* hidden,
    int N
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= N) return;

    float g = float(gate[idx]);
    float h = float(hidden[idx]);

    float grad_lc = float(grad_log_coeffs[idx]);
    float grad_lv = float(grad_log_values[idx]);
    float grad_g_from_lc = -softplus_bwd(grad_lc, g);
    float grad_g_from_lz = -softplus_bwd(-grad_lv, -g);
    float grad_g_total = grad_g_from_lc + grad_g_from_lz;
    grad_gate[idx] = T(grad_g_total);
    float log_tilde_h;
    float grad_h_from_lt;
    if (h >= 0.0f) {
        float relu_h = relu(h);
        log_tilde_h = logf(relu_h + 0.5f);
        float inner_grad = 1.0f / (relu_h + 0.5f);
        grad_h_from_lt = relu_backward(h, inner_grad * grad_lv);
    } else {
        log_tilde_h = -softplus_fwd(-h);
        grad_h_from_lt = -softplus_bwd(-grad_lv, -h);
    }
    grad_hidden[idx] = T(grad_h_from_lt);
}

template<typename T>
void launch_log_coeffs_and_values(
    T* log_coeffs,
    T* log_values,
    const T* gate,
    const T* hidden,
    int N,
    cudaStream_t stream
) {
    int grid = grid_size(N);
    log_coeffs_and_values_kernel<T><<<grid, BLOCK_SIZE, 0, stream>>>(
        log_coeffs,
        log_values,
        gate,
        hidden,
        N
    );

    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        fprintf(stderr, "CUDA kernel launch error: %s\n", cudaGetErrorString(err));
    }
}

template<typename T>
void launch_log_coeffs_and_values_backward(
    T* grad_gate,
    T* grad_hidden,
    const T* grad_log_coeffs,
    const T* grad_log_values,
    const T* gate,
    const T* hidden,
    int N,
    cudaStream_t stream
) {
    int grid = grid_size(N);
    log_coeffs_and_values_backward_kernel<T><<<grid, BLOCK_SIZE, 0, stream>>>(
        grad_gate,
        grad_hidden,
        grad_log_coeffs,
        grad_log_values,
        gate,
        hidden,
        N
    );

    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        fprintf(stderr, "CUDA kernel launch error: %s\n", cudaGetErrorString(err));
    }
}

__device__ __forceinline__ double logcumsumexp_forward(double x, double acc) {
    if (acc == -INFINITY) {
        return x;
    } else {
        double min_val = fmin(acc, x);
        double max_val = fmax(acc, x);
        return max_val + log1pf(expf(min_val - max_val));
    }
}

__device__ __forceinline__ double logcumsumexp_backward(double x, double* acc, double grad, double s, double* s_nxt) {
    *acc = grad + *acc * exp(s - *s_nxt);
    *s_nxt = s;
    return *acc * exp(x - s);
}

// float32 + branch free
__device__ __forceinline__ float logcumsumexp_forward_opt(float x, float acc) {
    float min_val = fminf(acc, x);
    float max_val = fmaxf(acc, x);
    return max_val + log1pf(__expf(min_val - max_val));
}

__device__ __forceinline__ float logcumsumexp_backward_opt(float x, float* acc, float grad, float s, float* s_nxt) {
    *acc = fmaf(*acc, __expf(s - *s_nxt), grad);
    *s_nxt = s;
    return *acc * __expf(x - s);
}

// Fully fused forward: chunk + log_coeffs_and_values + scan + sigmoid(proj)*out
// Takes combined (B, T, 3*H) = [hidden, gate, proj] and outputs gated result
template<typename T>
__global__ void fused_scan_forward_kernel(
    T* __restrict__ out,                 // (B, T, H) - sigmoid(proj) * scan_result
    T* __restrict__ next_state,          // (B, 1, H) - raw scan_result at T (for recurrence)
    float* __restrict__ a_star_buf,      // (B, T+1, H) - for backward
    float* __restrict__ s_buf,           // (B, T+1, H) - for backward
    float* __restrict__ log_values_buf,  // (B, T+1, H) - cached log_values for backward
    const T* __restrict__ combined,      // (B, T, 3*H) = [hidden(H), gate(H), proj(H)]
    const T* __restrict__ state,         // (B, 1, H)
    int T_seq,                           // sequence length (T)
    int H,
    int B
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= B * H) return;

    int b = idx / H;
    int h = idx % H;

    int T_out = T_seq + 1;
    int buf_base = b * T_out * H + h;    // base for a_star/s/log_values buffers (T+1 timesteps)
    int out_base = b * T_seq * H + h;    // base for output (T timesteps)
    int state_idx = b * H + h;           // state is (B, 1, H) -> flatten to (B, H)

    float a_star = 0.0f;
    float s = -INFINITY;  // logcumsumexp accumulator

    // Handle t=0 outside the loop: use log(state), coeff = 0
    float log_value_0 = logf(float(state[state_idx]));
    log_values_buf[buf_base] = log_value_0;
    s = log_value_0;  // z = log_value - a_star = log_value - 0 = log_value
    a_star_buf[buf_base] = a_star;
    s_buf[buf_base] = s;

    // Loop t=1..T_seq (no branches needed)
    float scan_result = 0.0f;
    for (int t = 1; t < T_out; t++) {
        int buf_curr = buf_base + t * H;
        int combined_base = b * T_seq * 3 * H + (t - 1) * 3 * H;

        float hidden_val = float(combined[combined_base + h]);
        float gate_val = float(combined[combined_base + H + h]);
        float proj_val = float(combined[combined_base + 2 * H + h]);

        float log_coeff_val, log_value_val;
        log_coeffs_and_values_fwd(gate_val, hidden_val, &log_coeff_val, &log_value_val);

        // Cache log_value for backward (avoid recomputation)
        log_values_buf[buf_curr] = log_value_val;

        // a_star[t] = sum_{i=0}^t log_coeffs[i]
        a_star += log_coeff_val;

        float z = log_value_val - a_star;

        if (s == -INFINITY) {
            s = z;
        } else {
            float min_val = fminf(s, z);
            float max_val = fmaxf(s, z);
            s = max_val + log1pf(expf(min_val - max_val));
        }

        scan_result = expf(a_star + s);

        // sigmoid(proj) * out
        int out_curr = out_base + (t - 1) * H;
        float proj_sigmoid = sigmoid(proj_val);
        out[out_curr] = T(proj_sigmoid * scan_result);

        a_star_buf[buf_curr] = a_star;
        s_buf[buf_curr] = s;
    }
    // Write timestep T to next_state (raw scan_result, no proj, for recurrence)
    next_state[state_idx] = T(scan_result);
}

// Fully fused backward: chains through sigmoid(proj)*out and log_coeffs_and_values
// Takes combined (B, T, 3*H), outputs grad_combined (B, T, 3*H) = [grad_hidden, grad_gate, grad_proj]
template<typename T>
__global__ void fused_scan_backward_kernel(
    T* __restrict__ grad_combined,         // (B, T, 3*H) = [grad_hidden, grad_gate, grad_proj]
    T* __restrict__ grad_state,            // (B, 1, H)
    const T* __restrict__ grad_out,        // (B, T, H) - gradient of sigmoid(proj)*scan_result
    const T* __restrict__ grad_next_state, // (B, 1, H) - gradient of raw scan_result at T
    const T* __restrict__ combined,        // (B, T, 3*H) = [hidden, gate, proj]
    const T* __restrict__ state,           // (B, 1, H)
    const float* __restrict__ a_star_buf,  // (B, T+1, H)
    const float* __restrict__ s_buf,       // (B, T+1, H)
    const float* __restrict__ log_values_buf, // (B, T+1, H) - cached from forward
    int T_seq,                             // sequence length (T)
    int H,
    int B
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= B * H) return;

    int b = idx / H;
    int h = idx % H;

    int T_out = T_seq + 1;
    int buf_base = b * T_out * H + h;    // base for a_star/s/log_values buffers (T+1 timesteps)
    int out_base = b * T_seq * H + h;    // base for grad_out (T timesteps)
    int state_idx = b * H + h;           // state is (B, 1, H) -> flatten to (B, H)

    float acc = 0.0;
    float s_val_next = 0.0;
    float carry_grad_a = 0.0;

    for (int t = T_out - 1; t >= 0; --t) {
        int base_adr = b*T_seq*3*H + (t-1)*3*H;
        int hidden_adr = base_adr + h;
        int gate_adr = base_adr + H + h;
        int proj_adr = base_adr + 2*H + h;

        int buf_curr = buf_base + t * H;

        float a_star = a_star_buf[buf_curr];
        float s = s_buf[buf_curr];
        float scan_result = expf(a_star + s);  // reconstruct scan result

        // Read cached log_value from forward pass (no recomputation needed)
        float log_value_val = log_values_buf[buf_curr];

        // Read from combined for t >= 1 (still need gate/hidden for backward, proj for output gate)
        float gate_val = 0.0f, hidden_val = 0.0f, proj_val = 0.0f;
        int combined_base = 0;

        if (t >= 1) {
            hidden_val = float(combined[hidden_adr]);
            gate_val = float(combined[gate_adr]);
            proj_val = float(combined[proj_adr]);
        }

        float z = log_value_val - a_star;

        // Get gradient for this timestep
        // For t >= 1: grad_out is gradient of (sigmoid(proj) * scan_result)
        // For t = T: also add grad_next_state (gradient of raw scan_result)
        float grad_gated_out = 0.0f;
        float grad_scan_from_next = 0.0f;

        if (t >= 1) {
            int grad_out_idx = out_base + (t - 1) * H;
            grad_gated_out = float(grad_out[grad_out_idx]);
        }
        if (t == T_seq) {
            grad_scan_from_next = float(grad_next_state[state_idx]);
        }

        // Chain through sigmoid(proj) * scan_result
        // out = sigmoid(proj) * scan_result
        // d_out/d_scan_result = sigmoid(proj)
        // d_out/d_proj = scan_result * sigmoid(proj) * (1 - sigmoid(proj))
        float grad_scan_result = grad_scan_from_next;
        float grad_proj = 0.0f;

        if (t >= 1) {
            float proj_sigmoid = sigmoid(proj_val);
            grad_scan_result += grad_gated_out * proj_sigmoid;
            // sigmoid'(x) = sigmoid(x) * (1 - sigmoid(x))
            grad_proj = grad_gated_out * scan_result * proj_sigmoid * (1.0f - proj_sigmoid);
        }

        // Now chain grad_scan_result through the scan backward
        float grad_log_h = grad_scan_result * scan_result;
        float grad_s = grad_log_h;

        if (t == T_out - 1) {
            acc = grad_s;
        } else {
            acc = grad_s + acc * expf(s - s_val_next);
        }
        float grad_z = acc * expf(z - s);
        s_val_next = s;

        float grad_a = grad_log_h + carry_grad_a - grad_z;
        carry_grad_a = grad_a;

        if (t == 0) {
            // grad_state = grad_z * d(log(state))/d(state) = grad_z / state
            grad_state[state_idx] = T(grad_z / float(state[state_idx]));
        } else {
            // Chain through log_coeffs_and_values backward to get grad_gate, grad_hidden
            float grad_g, grad_h;
            log_coeffs_and_values_bwd(grad_a, grad_z, gate_val, hidden_val, &grad_g, &grad_h);

            // Write to grad_combined: [grad_hidden, grad_gate, grad_proj]
            grad_combined[gate_adr] = T(grad_g);
            grad_combined[hidden_adr] = T(grad_h);
            grad_combined[proj_adr] = T(grad_proj);
        }
    }
}

/*
template<typename T>
__global__ void fused_scan_backward_kernel(
    T* __restrict__ grad_log_coeffs,
    T* __restrict__ grad_log_values,
    const T* __restrict__ grad_out,
    const T* __restrict__ out_buf,
    const double* __restrict__ a_star_buf,
    const double* __restrict__ s_buf,
    const T* __restrict__ log_values,
    int T_total,
    int H,
    int B
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= B * H) return;

    int b = idx / H;
    int h = idx % H;

    int base = b * T_total * H + h;

    double carry_grad_a = 0.0;
    double carry_grad_s = 0.0;

    for (int t = T_total - 1; t >= 0; --t) {
        int curr = base + t * H;

        double a_star = a_star_buf[curr];
        double s = s_buf[curr];
        double z = double(log_values[curr]) - a_star;
        double grad_log_h = double(grad_out[curr]) * double(out_buf[curr]); // out_buf[t] = exp(log_h[t])

        double grad_s = grad_log_h + carry_grad_s;

        double s_prev = -INFINITY;
        if (t > 0) {
            s_prev = s_buf[base + (t - 1) * H];
        }

        double max_val = fmax(s_prev, z);

        double exp_prev = 0.0;
        if (s_prev != -INFINITY) {
            exp_prev = exp(s_prev - max_val);
        }

        double exp_z = 0.0;
        if (z != -INFINITY) {
            exp_z = exp(z - max_val);
        }

        double denom = exp_prev + exp_z;

        double frac_prev = 0.0;
        double frac_z = 0.0;
        if (denom != 0.0) {
            frac_prev = exp_prev / denom;
            frac_z = exp_z / denom;
        }

        // grad_z = (grad_log_h + carry_grad_s) * exp(z - max_val) / (exp(s_prev - max_val) + exp(z - max_val))
        // grad_z = (grad_log_h + exp(s - exp_nxt)) * exp(z - s) 

        double d_Z = frac_z * grad_s;
        double d_A = grad_log_h + carry_grad_a - d_Z;

        grad_log_values[curr] = T(d_Z);
        grad_log_coeffs[curr] = T(d_A);

        carry_grad_a = d_A;
        carry_grad_s = frac_prev * grad_s;
    }
}
*/


/*
template<typename T>
__global__ void fused_scan_backward_kernel(
    T* __restrict__ grad_log_coeffs,
    T* __restrict__ grad_log_values,
    const T* __restrict__ grad_out,
    const T* __restrict__ log_coeffs,
    const T* __restrict__ log_values,
    const T* __restrict__ out,
    const double* __restrict__ a_star_buf,
    const double* __restrict__ s_buf,
    int T_total,
    int H,
    int B
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= B * H) return;

    int b = idx / H;
    int h = idx % H;

    int base = b * T_total * H + h;

    double grad_a_star[1025] = {0};  // Assuming T_total <= 1024
    double W = 0.0;  // Accumulates sum_{i=t}^{T-1} [grad_log_h[i] * exp(-s[i])]

    for (int t = T_total - 1; t >= 0; t--) {
        int curr = base + t * H;

        double a_star = a_star_buf[curr];
        double s_val = s_buf[curr];
        double z_val = double(log_values[curr]) - a_star;

        // Compute dL/d(log_h[t]) = dL/d(out[t]) * d(out[t])/d(log_h[t])
        double grad_log_h = double(grad_out[curr]) * double(out[curr]);

        // Update W: W[t] = grad_log_h[t] * exp(-s_val) + W[t+1]
        W = grad_log_h * exp(-s_val) + W;

        // Compute dL/d(z[t]) = exp(z_val) * W[t]
        double grad_z = exp(z_val) * W;

        // dL/d(log_values[t]) = dL/d(z[t]) * dz[t]/d(log_values[t]) = grad_z
        grad_log_values[curr] = T(grad_z);

        // dL/da_star[t] = dL/d(log_h[t]) - dL/d(z[t]) (due to chain rule)
        grad_a_star[t] = grad_log_h - grad_z;
    }

    // Compute dL/d(log_coeffs) via cumulative sum of dL/da_star
    double accum = 0.0;
    for (int t = T_total - 1; t >= 0; t--) {
        accum += grad_a_star[t];
        grad_log_coeffs[base + t * H] = T(accum);
    }
}
*/


/*
template<typename T>
__global__ void fused_scan_backward_kernel(
    T* __restrict__ grad_log_coeffs,
    T* __restrict__ grad_log_values,
    const T* __restrict__ grad_out,
    const T* __restrict__ log_coeffs,
    const T* __restrict__ log_values,
    const T* __restrict__ out,
    const float* __restrict__ a_star_buf,
    const float* __restrict__ s_buf,
    int T_total,
    int H,
    int B
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= B * H) return;

    int b = idx / H;
    int h = idx % H;

    int base = b * T_total * H + h;

    float grad_a_star[1025] = {0};
    float G = 0.0f;  // G[t] = sum_{i=t}^{T-1} grad_s[i]
    for (int t = T_total - 1; t >= 0; t--) {
        int curr = base + t * H;

        float a_star = a_star_buf[curr];
        float s_val = s_buf[curr];
        float z = float(log_values[curr]) - a_star;

        // grad_log_h[t] = grad_out[t] * out[t]
        float grad_log_h = float(grad_out[curr]) * float(out[curr]);

        // G = sum of grad_s from t to end (grad_s[t] = grad_log_h[t])
        G += grad_log_h;

        // grad_z[t] = exp(z - s_val) * G
        float prob = expf(z - s_val);
        float grad_z = prob * G;

        // grad_log_values[t] = grad_z
        grad_log_values[curr] = T(grad_z);

        // grad_a_star[t] gets:
        // - +grad_log_h (from log_h = a_star + s)
        // - -grad_z    (from z = log_values - a_star)
        grad_a_star[t] = grad_log_h - grad_z;
    }

    // grad_log_coeffs[t] = sum_{i=t}^{T-1} grad_a_star[i]
    float accum = 0.0f;
    for (int t = T_total - 1; t >= 0; t--) {
        accum += grad_a_star[t];
        grad_log_coeffs[base + t * H] = T(accum);
    }
}

 template<typename T>
__global__ void fused_scan_backward_kernel(
    T* __restrict__ grad_log_coeffs,
    T* __restrict__ grad_log_values,
    const T* __restrict__ grad_out,
    const T* __restrict__ log_coeffs,
    const T* __restrict__ log_values,
    const T* __restrict__ out,
    const float* __restrict__ a_star_buf,
    const float* __restrict__ s_buf,
    int T_total,
    int H,
    int B
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= B * H) return;

    int b = idx / H;
    int h = idx % H;

    int base = b * T_total * H + h;

    // Recompute z[t] = log_values[t] - a_star[t]
    float z[1025];
    for (int t = 0; t < T_total; t++) {
        int curr = base + t * H;
        z[t] = float(log_values[curr]) - a_star_buf[curr];
    }

    // g_log_h[t] = grad_out[t] * out[t]
    float g_log_h[1025];
    for (int t = 0; t < T_total; t++) {
        int curr = base + t * H;
        g_log_h[t] = float(grad_out[curr]) * float(out[curr]);
    }

    // Step: Online logcumsumexp backward for g_z
    float g_z[1025] = {0};
    g_z[T_total - 1] = g_log_h[T_total - 1];

    for (int t = T_total - 2; t >= 0; t--) {
        float exp_term = expf(z[t] - s_buf[base + (t + 1) * H]);
        g_z[t] = g_log_h[t] + g_z[t + 1] * exp_term;
    }

    // grad_log_values[t] = g_z[t]
    for (int t = 0; t < T_total; t++) {
        int curr = base + t * H;
        grad_log_values[curr] = T(g_z[t]);
    }

    // g_a_star[t] = g_log_h[t] - g_z[t]
    float g_a_star[1025] = {0};
    for (int t = 0; t < T_total; t++) {
        g_a_star[t] = g_log_h[t] - g_z[t];
    }

    // grad_log_coeffs[t] = reverse cumsum of g_a_star
    float accum = 0.0f;
    for (int t = T_total - 1; t >= 0; t--) {
        accum += g_a_star[t];
        grad_log_coeffs[base + t * H] = T(accum);
    }
}
*/
// Fully fused forward launch: takes combined (B, T, 3*H) = [hidden, gate, proj]
template<typename T>
void launch_fused_scan_forward(
    T* out,
    T* next_state,
    float* a_star,
    float* s_vals,
    float* log_values_buf,  // (B, T+1, H) - cached for backward
    const T* combined,  // (B, T, 3*H) = [hidden, gate, proj]
    const T* state,
    int T_seq,
    int H,
    int B,
    cudaStream_t stream
) {
    int total = B * H;
    int grid = seq_size(total);

    fused_scan_forward_kernel<T><<<grid, SEQ_SIZE, 0, stream>>>(
        out,
        next_state,
        a_star,
        s_vals,
        log_values_buf,
        combined,
        state,
        T_seq,
        H,
        B
    );

    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        fprintf(stderr, "CUDA kernel launch error in forward: %s\n", cudaGetErrorString(err));
    }
}

// Fully fused backward launch: outputs grad_combined (B, T, 3*H) = [grad_hidden, grad_gate, grad_proj]
template<typename T>
void launch_fused_scan_backward(
    T* grad_combined,   // (B, T, 3*H) = [grad_hidden, grad_gate, grad_proj]
    T* grad_state,
    const T* grad_out,
    const T* grad_next_state,
    const T* combined,  // (B, T, 3*H) = [hidden, gate, proj]
    const T* state,
    const float* a_star_buf,
    const float* s_buf,
    const float* log_values_buf,  // (B, T+1, H) - cached from forward
    int T_seq,
    int H,
    int B,
    cudaStream_t stream
) {
    int total = B * H;
    int grid = seq_size(total);

    fused_scan_backward_kernel<T><<<grid, SEQ_SIZE, 0, stream>>>(
        grad_combined,
        grad_state,
        grad_out,
        grad_next_state,
        combined,
        state,
        a_star_buf,
        s_buf,
        log_values_buf,
        T_seq,
        H,
        B
    );

    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        fprintf(stderr, "CUDA kernel launch error in backward: %s\n", cudaGetErrorString(err));
    }
}

/*
__device__ __forceinline__ float log_add_exp(const float a, const float b) {
  if (::isnan(a) || ::isnan(b)) {
    return std::numeric_limits<float>::quiet_NaN();
  }
  float min_val = fminf(a, b);
  float max_val = fmaxf(a, b);
  if (min_val != max_val || ::isfinite(min_val)) {
    return max_val + log1pf(expf(min_val - max_val));
  } else {
      return a;
  }
}

__device__ __forceinline__ float log_add_exp_backward(float x_val, float s_val) {
  if (::isnan(x_val) || ::isnan(s_val)) {
    return std::numeric_limits<float>::quiet_NaN();
  }
  return expf(x_val - s_val);
}
*/

__device__ __forceinline__ double log_add_exp(const double a, const double b) {
  double min_val = fmin(a, b);
  double max_val = fmax(a, b);
  return max_val + log1p(exp(min_val - max_val));
}

__device__ __forceinline__ double log_add_exp_backward(double x, double s) {
    return exp(x - s);
}

 
// This exactly matches pytorch in double, but not in float
template<typename T>
__global__ void logcumsumexp_forward_kernel(
    T* __restrict__ out,           // exp(s[t])
    double* __restrict__ s_buf,     // s[t] = logcumsumexp(x[0..t])
    const T* __restrict__ x,       // input: log_values
    int T_total,
    int H,
    int B
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= B * H) return;

    int b = idx / H;
    int h = idx % H;

    int base = b * T_total * H + h;

    double s = -INFINITY;

    for (int t = 0; t < T_total; t++) {
        int curr = base + t * H;
        double x_val = double(x[curr]);
        s = logcumsumexp_forward(x_val, s);
        out[curr] = T(s);
        s_buf[curr] = s;
    }
}
template<typename T>
__global__ void logcumsumexp_backward_kernel(
    T* __restrict__ grad_x,
    const T* __restrict__ grad_out,
    const T* __restrict__ x,
    const double* __restrict__ s_buf,
    int T_total,
    int H,
    int B
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= B * H) return;

    int b = idx / H;
    int h = idx % H;

    int base = b * T_total * H + h;

    double acc = 0.0;
    double s_val_next = 0.0;

    for (int t = T_total - 1; t >= 0; --t) {
        int curr = base + t * H;

        double x_val = double(x[curr]);
        double s_val = double(s_buf[curr]);
        double g_val = double(grad_out[curr]);
        grad_x[curr] = T(logcumsumexp_backward(x_val, &acc, g_val, s_val, &s_val_next));
    }
}
/*
template<typename T>
__global__ void logcumsumexp_backward_kernel(
    T* __restrict__ grad_x,
    const T* __restrict__ grad_out,
    const T* __restrict__ x,
    const float* __restrict__ s_buf,
    int T_total,
    int H,
    int B
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= B * H) return;

    int b = idx / H;
    int h = idx % H;

    int base = b * T_total * H + h;

    // grad_x[i] = sum_{t≥i} grad_out[t] * exp(x[i] - s[t])
    for (int i = 0; i < T_total; i++) {
        int curr_i = base + i * H;
        float x_i = float(x[curr_i]);
        float g = 0.0f;

        for (int t = i; t < T_total; t++) {
            int curr_t = base + t * H;
            float s_t = s_buf[curr_t];
            float prob = expf(x_i - s_t);
            //float prob = log_add_exp_backward(x_i, s_t);
            g += float(grad_out[curr_t]) * prob;
        }

        grad_x[curr_i] = T(g);
    }
}
*/

template<typename T>
void launch_logcumsumexp_forward(
    T* out,
    double* s_buf,
    const T* x,
    int T_total,
    int H,
    int B,
    cudaStream_t stream
) {
    int total = B * H;
    int grid = grid_size(total);

    logcumsumexp_forward_kernel<T><<<grid, BLOCK_SIZE, 0, stream>>>(
        out, s_buf, x, T_total, H, B
    );

    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess)
        fprintf(stderr, "Forward kernel error: %s\n", cudaGetErrorString(err));
}

template<typename T>
void launch_logcumsumexp_backward(
    T* grad_x,
    const T* grad_out,
    const T* x,
    const double* s_buf,
    int T_total,
    int H,
    int B,
    cudaStream_t stream
) {
    int total = B * H;
    int grid = grid_size(total);

    logcumsumexp_backward_kernel<T><<<grid, BLOCK_SIZE, 0, stream>>>(
        grad_x, grad_out, x, s_buf, T_total, H, B
    );

    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess)
        fprintf(stderr, "Backward kernel error: %s\n", cudaGetErrorString(err));
}

// logcumsumexp in float32
template<typename T>
__global__ void logcumsumexp_forward_kernel_opt(
    T* __restrict__ out,
    float* __restrict__ s_buf,  // FLOAT32
    const T* __restrict__ x,
    int T_total,
    int H,
    int B
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= B * H) return;

    int b = idx / H;
    int h = idx % H;

    int base = b * T_total * H + h;

    float s = -INFINITY;

    for (int t = 0; t < T_total; t++) {
        int curr = base + t * H;
        float x_val = float(x[curr]);
        s = logcumsumexp_forward_opt(x_val, s);
        out[curr] = T(s);
        s_buf[curr] = s;
    }
}

template<typename T>
__global__ void logcumsumexp_backward_kernel_opt(
    T* __restrict__ grad_x,
    const T* __restrict__ grad_out,
    const T* __restrict__ x,
    const float* __restrict__ s_buf,
    int T_total,
    int H,
    int B
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= B * H) return;

    int b = idx / H;
    int h = idx % H;

    int base = b * T_total * H + h;

    float acc = 0.0f;
    float s_val_next = 0.0f;

    for (int t = T_total - 1; t >= 0; --t) {
        int curr = base + t * H;

        float x_val = float(x[curr]);
        float s_val = s_buf[curr];
        float g_val = float(grad_out[curr]);
        grad_x[curr] = T(logcumsumexp_backward_opt(x_val, &acc, g_val, s_val, &s_val_next));
    }
}

template<typename T>
void launch_logcumsumexp_forward_opt(
    T* out,
    float* s_buf,
    const T* x,
    int T_total,
    int H,
    int B,
    cudaStream_t stream
) {
    int total = B * H;
    int grid = grid_size(total);

    logcumsumexp_forward_kernel_opt<T><<<grid, BLOCK_SIZE, 0, stream>>>(
        out, s_buf, x, T_total, H, B
    );

    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess)
        fprintf(stderr, "Forward opt kernel error: %s\n", cudaGetErrorString(err));
}

template<typename T>
void launch_logcumsumexp_backward_opt(
    T* grad_x,
    const T* grad_out,
    const T* x,
    const float* s_buf,
    int T_total,
    int H,
    int B,
    cudaStream_t stream
) {
    int total = B * H;
    int grid = grid_size(total);

    logcumsumexp_backward_kernel_opt<T><<<grid, BLOCK_SIZE, 0, stream>>>(
        grad_x, grad_out, x, s_buf, T_total, H, B
    );

    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess)
        fprintf(stderr, "Backward opt kernel error: %s\n", cudaGetErrorString(err));
}

template<typename T>
__global__ void ppo_loss_forward_kernel(
    float* __restrict__ loss,
    double* __restrict__ saved_for_backward,
    const T* __restrict__ logits,
    const T* __restrict__ values_pred,
    const int64_t* __restrict__ actions,
    const T* __restrict__ old_logprobs,
    const T* __restrict__ advantages,
    const T* __restrict__ prio,
    const T* __restrict__ values,
    const T* __restrict__ returns,
    const float* __restrict__ adv_mean,
    const float* __restrict__ adv_std,
    double clip_coef,
    double vf_clip_coef,
    double vf_coef,
    double ent_coef,
    int T_seq,
    int A,
    int N
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total_elements = N * T_seq;
    if (idx >= total_elements) return;
    __shared__ float block_loss[BLOCK_SIZE];

    int n = idx / T_seq;  // batch index
    int t = idx % T_seq;  // timestep

    // === Direct indexing: no lambdas ===
    int nt = n * T_seq + t;                    // index into (N, T_seq) tensors
    int logits_offset = n * T_seq * A + t * A; // base index into logits

    // === Step 1: Read action and compute logsumexp ===
    int act = actions[nt];  // action taken at (n,t)

    // Compute logsumexp: log(sum_a exp(logits[a]))
    double max_logit = -INFINITY;
    for (int a = 0; a < A; a++) {
        double l = double(logits[logits_offset + a]);
        max_logit = fmax(max_logit, l);
    }

    double logsumexp = 0.0;
    double sum = 0.0;
    for (int a = 0; a < A; a++) {
        double l = double(logits[logits_offset + a]);
        sum += exp(l - max_logit);
    }
    logsumexp = max_logit + log(sum);

    // === Step 2: new_logprob[action] = logits[action] - logsumexp ===
    // log_softmax = (logits - max_logit) - max_logit - logsumexp
    double new_logp = double(logits[logits_offset + act]) - logsumexp;

    // === Step 3: entropy = -sum_a p_a * log p_a ===
    double entropy = 0.0;
    for (int a = 0; a < A; a++) {
        double l = double(logits[logits_offset + a]);
        double p = exp(l - logsumexp);
        double logp = l - logsumexp;
        entropy -= p * logp;
    }

    // === Step 4: policy gradient loss ===
    double old_logp = double(old_logprobs[nt]);
    double adv = double(advantages[nt]);
    double w = double(prio[n]);  // importance weight, per-sequence
    double adv_normalized = (adv - adv_mean[0]) / (adv_std[0] + 1e-8);

    double logratio = new_logp - old_logp;
    double ratio = exp(logratio);

    double ratio_clipped = fmax(1.0 - clip_coef, fmin(1.0 + clip_coef, ratio));
    double pg_loss1 = -w * adv_normalized * ratio;
    double pg_loss2 = -w * adv_normalized * ratio_clipped;
    double pg_loss = fmax(pg_loss1, pg_loss2);  // PPO clipped surrogate loss

    // === Step 5: value function loss ===
    double val = double(values[nt]);
    double ret = double(returns[nt]);
    double val_pred = double(values_pred[nt]);

    double v_error = val_pred - val;
    double v_clipped = val + fmax(-vf_clip_coef, fmin(vf_clip_coef, v_error));
    double v_loss_unclipped = (val_pred - ret) * (val_pred - ret);
    double v_loss_clipped = (v_clipped - ret) * (v_clipped - ret);
    double v_loss = 0.5f * fmax(v_loss_unclipped, v_loss_clipped);

    // === Step 6: total sample loss (pre-divided by N*T for mean) ===
    double thread_loss = (pg_loss + vf_coef * v_loss - ent_coef * entropy) / double(total_elements);

    // === Save for backward ===
    double* saved_row = saved_for_backward + idx * 5;
    saved_row[0] = new_logp;
    saved_row[1] = ratio;
    saved_row[2] = val_pred;
    saved_row[3] = v_clipped;
    saved_row[4] = entropy;

    // === Block-local reduction using shared memory ===
    int tid = threadIdx.x;
    block_loss[tid] = float(thread_loss);
    __syncthreads();

    // Reduce within block using tree reduction
    for (int stride = BLOCK_SIZE / 2; stride > 0; stride >>= 1) {
        if (tid < stride) {
            block_loss[tid] += block_loss[tid + stride];
        }
        __syncthreads();
    }

    // === Accumulate into loss_output (scalar via atomic add) ===
    if (tid == 0) {
        atomicAdd(loss, block_loss[0]);
    }
}

template<typename T>
__global__ void ppo_loss_backward_kernel(
    T* __restrict__ grad_logits,
    T* __restrict__ grad_values_pred,
    const float* __restrict__ grad_loss,  // scalar, [1], dL/dloss
    const T* __restrict__ logits,
    const int64_t* __restrict__ actions,
    const T* __restrict__ old_logprobs,
    const T* __restrict__ advantages,
    const T* __restrict__ prio,
    const T* __restrict__ values,
    const T* __restrict__ returns,
    const double* __restrict__ saved_for_backward,
    const float* __restrict__ adv_mean,
    const float* __restrict__ adv_std,
    double clip_coef,
    double vf_clip_coef,
    double vf_coef,
    double ent_coef,
    int T_seq,
    int A,
    int N
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total_elements = N * T_seq;
    if (idx >= total_elements) return;

    double inv_NT = 1.0f / (N * T_seq);
    int n = idx / T_seq;
    int t = idx % T_seq;

    // === Direct indexing ===
    int nt = n * T_seq + t;
    int logits_offset = n * T_seq * A + t * A;

    // === Retrieve saved values from forward pass ===
    const double* saved = saved_for_backward + idx * 5;
    double new_logp = saved[0];   // new log prob of selected action
    double ratio = saved[1];      // exp(new_logp - old_logp)
    double val_pred = saved[2];   // value prediction
    double v_clipped = saved[3];  // clipped value target
    double entropy = saved[4];    // entropy at (n,t)

    // === Read inputs ===
    double old_logp = double(old_logprobs[nt]);
    double adv = double(advantages[nt]);
    double w = double(prio[n]);  // importance weight
    double val = double(values[nt]);
    double ret = double(returns[nt]);

    // === Normalize advantage (same as forward) ===
    double adv_normalized = (adv - adv_mean[0]) / (adv_std[0] + 1e-8f);

    // Total loss gradient (scalar from autograd)
    double dL = grad_loss[0] * inv_NT;  // dL/dloss

    // Gradients w.r.t. components
    double d_pg_loss = dL;                    // policy loss contributes dL
    double d_v_loss = dL * vf_coef;           // value loss scaled by vf_coef
    double d_entropy_term = dL * (-ent_coef); // entropy bonus gradient

    // ===================================================
    // 1. Gradient w.r.t. value function prediction
    // ===================================================
    double v_loss_unclipped = (val_pred - ret) * (val_pred - ret);
    double v_loss_clipped = (v_clipped - ret) * (v_clipped - ret);

    // Which branch was taken in forward? (same logic as PyTorch: use unclipped if tie)
    bool use_clipped_vf = (v_loss_clipped > v_loss_unclipped);
    double d_val_pred = 0.0;

    if (use_clipped_vf) {
        double v_error = val_pred - val;
        if (v_error >= -vf_clip_coef && v_error <= vf_clip_coef) {
            d_val_pred = v_clipped - ret;  // = val_pred - ret
        }
    } else {
        d_val_pred = val_pred - ret;
    }

    d_val_pred = dL * vf_coef * d_val_pred;
    grad_values_pred[nt] = T(d_val_pred);

    // ===================================================
    // 2. Gradient w.r.t. policy and entropy (logits)
    // ===================================================
    // Recompute logsumexp for gradient
    double max_logit = -INFINITY;
    for (int a = 0; a < A; a++) {
        double l = double(logits[logits_offset + a]);
        max_logit = fmax(max_logit, l);
    }

    double logsumexp = 0.0;
    double sum = 0.0;
    for (int a = 0; a < A; a++) {
        double l = double(logits[logits_offset + a]);
        sum += exp(l - max_logit);
    }
    logsumexp = max_logit + log(sum);
 
    // Zero grad_logits for this (n,t)
    for (int a = 0; a < A; a++) {
        grad_logits[logits_offset + a] = T(0.0f);
    }

    // --- Policy Loss Gradient ---
    double logratio = new_logp - old_logp;
    double ratio_clipped = fmax(1.0f - clip_coef, fmin(1.0f + clip_coef, ratio));
    double pg_loss1 = -w * adv_normalized * ratio;
    double pg_loss2 = -w * adv_normalized * ratio_clipped;

    double d_ratio = -w * adv_normalized * d_pg_loss;
    if (pg_loss2 > pg_loss1) {
        if (ratio <= (1.0 - clip_coef) || ratio >= (1.0 + clip_coef)) {
            d_ratio = 0.0;
        }
    }

    // d(ratio)/d(new_logp) = ratio
    double d_new_logp = d_ratio * ratio;

    // --- Entropy Gradient ---
    // dH/dlogits[a] = p_a * (entropy - log p_a)
    for (int a = 0; a < A; a++) {
        double l = double(logits[logits_offset + a]);
        double p = exp(l - logsumexp);
        double logp = l - logsumexp;

        // Gradient from policy loss: d/dlogits[a] new_logp = δ_{a,act} - p_a
        double d_logit = 0.0f;
        if (a == actions[nt]) {
            d_logit += d_new_logp;
        }
        d_logit -= p * d_new_logp;

        // Gradient from entropy
        // TODO: Grad is a bit more off than I would like (1e-6)
        // Probably need to check logsumexp (not cumulative) vs
        // torch / actually look at the puffer 3 entropy impl
        double d_entropy_dlogit = p * (entropy - logp);
        d_logit += d_entropy_term * d_entropy_dlogit;

        grad_logits[logits_offset + a] = T(d_logit);
    }
}

template<typename T>
inline void launch_ppo_loss_forward(
    float* loss_output,
    double* saved_for_backward,
    const T* logits,
    const T* values_pred,
    const int64_t* actions,
    const T* old_logprobs,
    const T* advantages,
    const T* prio,
    const T* values,
    const T* returns,
    const float* adv_mean,
    const float* adv_std,
    double clip_coef,
    double vf_clip_coef,
    double vf_coef,
    double ent_coef,
    int T_seq,
    int A,
    int N,
    cudaStream_t stream
) {
    int total_elements = N * T_seq;
    int grid = grid_size(total_elements);

    ppo_loss_forward_kernel<T><<<grid, BLOCK_SIZE, 0, stream>>>(
        loss_output,
        saved_for_backward,
        logits,
        values_pred,
        actions,
        old_logprobs,
        advantages,
        prio,
        values,
        returns,
        adv_mean,
        adv_std,
        clip_coef,
        vf_clip_coef,
        vf_coef,
        ent_coef,
        T_seq,
        A,
        N
    );

    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        fprintf(stderr, "PPO forward kernel error: %s\n", cudaGetErrorString(err));
    }
}

template<typename T>
void launch_ppo_loss_backward(
    T* grad_logits,
    T* grad_values_pred,
    const float* grad_loss,
    const T* logits,
    const int64_t* actions,
    const T* old_logprobs,
    const T* advantages,
    const T* prio,
    const T* values,
    const T* returns,
    const double* saved_for_backward,
    const float* adv_mean,
    const float* adv_std,
    double clip_coef,
    double vf_clip_coef,
    double vf_coef,
    double ent_coef,
    int T_seq,
    int A,
    int N,
    cudaStream_t stream
) {
    int total_elements = N * T_seq;
    int grid = grid_size(total_elements);

    ppo_loss_backward_kernel<T><<<grid, BLOCK_SIZE, 0, stream>>>(
        grad_logits,
        grad_values_pred,
        grad_loss,
        logits,
        actions,
        old_logprobs,
        advantages,
        prio,
        values,
        returns,
        saved_for_backward,
        adv_mean,
        adv_std,
        clip_coef,
        vf_clip_coef,
        vf_coef,
        ent_coef,
        T_seq,
        A,
        N
    );

    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        fprintf(stderr, "PPO backward kernel error: %s\n", cudaGetErrorString(err));
    }
}

// ============================================================================
// Fused sample_logits kernel: nan_to_num + log_softmax + multinomial + gather + value copy
// Inference-only (no gradients needed)
// Uses inline cuRAND to avoid separate torch::rand() kernel launch
//
// NOTE: This kernel supports strided (non-contiguous) logits/value input.
// This is needed for fused logit+value decoder output where logits is a view
// of a larger (B, V+A) tensor. The stride parameters handle this case
// to avoid .contiguous() kernel launches.
// ============================================================================

// Single kernel that handles: nan_to_num, log_softmax, multinomial sampling, logprob gather, value copy
// Input: logits (B, A) with row stride logits_stride, value (B, 1) with row stride value_stride, seed for RNG
// Output: actions (B,) as float64, logprobs (B,), value_out (B,)
// NOTE: offset is read from a pointer (not passed by value) so it works correctly with CUDA graphs.
// The offset tensor is incremented with a CUDA op after this kernel, so each graph replay gets a new offset.
template<typename T>
__global__ void sample_logits_kernel(
    double* __restrict__ actions,         // (B,) output - sampled action indices as float64
    T* __restrict__ logprobs,             // (B,) output - log prob of sampled action
    T* __restrict__ value_out,            // (B,) output - copied value (flattened)
    const T* __restrict__ logits,         // (B, A) input - raw logits (may be non-contiguous)
    const T* __restrict__ value,          // (B, 1) input - value from fused output (may be non-contiguous)
    uint64_t seed,                        // RNG seed
    const int64_t* __restrict__ offset_ptr, // RNG offset pointer (read at execution time for CUDA graph support)
    int A,                                // number of actions
    int B,                                // batch size
    int logits_stride,                    // stride between rows (for non-contiguous logits from fused output)
    int value_stride                      // stride between rows (for non-contiguous value from fused output)
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= B) return;

    // Read offset at execution time (important for CUDA graph replay)
    uint64_t offset = static_cast<uint64_t>(*offset_ptr);

    int logits_base = idx * logits_stride;

    // Step 1: Find max for numerical stability (with nan_to_num)
    float max_val = -INFINITY;
    for (int a = 0; a < A; ++a) {
        float l = float(logits[logits_base + a]);
        // nan_to_num: replace nan with 0
        if (isnan(l)) l = 0.0f;
        // clamp inf/-inf (pytorch defaults: neginf=-3.4028e+38, posinf=3.4028e+38)
        if (isinf(l)) l = (l > 0) ? 3.4028e+38f : -3.4028e+38f;
        max_val = fmaxf(max_val, l);
    }

    // Step 2: Compute logsumexp for log_softmax denominator
    float sum_exp = 0.0f;
    for (int a = 0; a < A; ++a) {
        float l = float(logits[logits_base + a]);
        if (isnan(l)) l = 0.0f;
        if (isinf(l)) l = (l > 0) ? 3.4028e+38f : -3.4028e+38f;
        sum_exp += expf(l - max_val);
    }
    float logsumexp = max_val + logf(sum_exp);

    // Step 3: Generate random value using Philox RNG (fast, high quality)
    curandStatePhilox4_32_10_t state;
    curand_init(seed, idx, offset, &state);
    float rand_val = curand_uniform(&state);

    // Step 4: Multinomial sampling using inverse CDF
    float cumsum = 0.0f;
    int sampled_action = A - 1;  // default to last action

    for (int a = 0; a < A; ++a) {
        float l = float(logits[logits_base + a]);
        if (isnan(l)) l = 0.0f;
        if (isinf(l)) l = (l > 0) ? 3.4028e+38f : -3.4028e+38f;
        float prob = expf(l - logsumexp);
        cumsum += prob;
        if (rand_val < cumsum) {
            sampled_action = a;
            break;
        }
    }

    // Step 5: Gather log probability of sampled action
    float sampled_logit = float(logits[logits_base + sampled_action]);
    if (isnan(sampled_logit)) sampled_logit = 0.0f;
    if (isinf(sampled_logit)) sampled_logit = (sampled_logit > 0) ? 3.4028e+38f : -3.4028e+38f;
    float log_prob = sampled_logit - logsumexp;

    // Write outputs (action as float64 for compatibility with continuous/discrete)
    actions[idx] = double(sampled_action);
    logprobs[idx] = T(log_prob);

    // Copy value (fused to avoid separate elementwise kernel for strided->contiguous copy)
    value_out[idx] = value[idx * value_stride];

    // Increment RNG offset for next call (thread 0 only, fused to avoid separate kernel)
    if (idx == 0) {
        atomicAdd((unsigned long long*)offset_ptr, 1ULL);
    }
}

template<typename T>
void launch_sample_logits(
    double* actions,
    T* logprobs,
    T* value_out,
    const T* logits,
    const T* value,
    uint64_t seed,
    const int64_t* offset_ptr,  // pointer to offset tensor (read at execution time for CUDA graphs)
    int A,
    int B,
    int logits_stride,  // stride between rows (for non-contiguous logits from fused output)
    int value_stride,   // stride between rows (for non-contiguous value from fused output)
    cudaStream_t stream
) {
    int grid = grid_size(B);
    sample_logits_kernel<T><<<grid, BLOCK_SIZE, 0, stream>>>(
        actions,
        logprobs,
        value_out,
        logits,
        value,
        seed,
        offset_ptr,
        A,
        B,
        logits_stride,
        value_stride
    );

    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        fprintf(stderr, "sample_logits kernel error: %s\n", cudaGetErrorString(err));
    }
}

// Non-templated wrappers for float
void launch_mingru_gate_inference_float(float* out, float* next_state, const float* combined, const float* state_in, int H, int B, cudaStream_t stream) {
    launch_mingru_gate_inference<float>(out, next_state, combined, state_in, H, B, stream);
}
void launch_log_coeffs_and_values_float(float* log_coeffs, float* log_values, const float* gate, const float* hidden, int N, cudaStream_t stream) {
    launch_log_coeffs_and_values<float>(log_coeffs, log_values, gate, hidden, N, stream);
}
void launch_log_coeffs_and_values_backward_float(float* grad_gate, float* grad_hidden, const float* grad_log_coeffs, const float* grad_log_values, const float* gate, const float* hidden, int N, cudaStream_t stream) {
    launch_log_coeffs_and_values_backward<float>(grad_gate, grad_hidden, grad_log_coeffs, grad_log_values, gate, hidden, N, stream);
}
void launch_rmsnorm_forward_float(float* out, float* inv_norm_buf, const float* x, const float* weight, double eps, int T_total, int H, int B, cudaStream_t stream) {
    launch_rmsnorm_forward<float>(out, inv_norm_buf, x, weight, eps, T_total, H, B, stream);
}
void launch_rmsnorm_backward_float(float* grad_x, float* grad_weight, const float* grad_out, const float* inv_norm_buf, const float* x_buf, const float* weight, double eps, int T_total, int H, int B, cudaStream_t stream) {
    launch_rmsnorm_backward<float>(grad_x, grad_weight, grad_out, inv_norm_buf, x_buf, weight, eps, T_total, H, B, stream);
}
void launch_fused_scan_forward_float(float* out, float* next_state, float* a_star, float* s_vals, float* log_values_buf, const float* combined, const float* state, int T_seq, int H, int B, cudaStream_t stream) {
    launch_fused_scan_forward<float>(out, next_state, a_star, s_vals, log_values_buf, combined, state, T_seq, H, B, stream);
}
void launch_fused_scan_backward_float(float* grad_combined, float* grad_state, const float* grad_out, const float* grad_next_state, const float* combined, const float* state, const float* a_star_buf, const float* s_buf, const float* log_values_buf, int T_seq, int H, int B, cudaStream_t stream) {
    launch_fused_scan_backward<float>(grad_combined, grad_state, grad_out, grad_next_state, combined, state, a_star_buf, s_buf, log_values_buf, T_seq, H, B, stream);
}
void launch_logcumsumexp_forward_float(float* out, double* s_buf, const float* x, int T_total, int H, int B, cudaStream_t stream) {
    launch_logcumsumexp_forward<float>(out, s_buf, x, T_total, H, B, stream);
}
void launch_logcumsumexp_backward_float(float* grad_x, const float* grad_out, const float* x, const double* s_buf, int T_total, int H, int B, cudaStream_t stream) {
    launch_logcumsumexp_backward<float>(grad_x, grad_out, x, s_buf, T_total, H, B, stream);
}
void launch_logcumsumexp_forward_opt_float(float* out, float* s_buf, const float* x, int T_total, int H, int B, cudaStream_t stream) {
    launch_logcumsumexp_forward_opt<float>(out, s_buf, x, T_total, H, B, stream);
}
void launch_logcumsumexp_backward_opt_float(float* grad_x, const float* grad_out, const float* x, const float* s_buf, int T_total, int H, int B, cudaStream_t stream) {
    launch_logcumsumexp_backward_opt<float>(grad_x, grad_out, x, s_buf, T_total, H, B, stream);
}
void launch_ppo_loss_forward_float(float* loss_output, double* saved_for_backward, const float* logits, const float* values_pred, const int64_t* actions, const float* old_logprobs, const float* advantages, const float* prio, const float* values, const float* returns, const float* adv_mean, const float* adv_std, double clip_coef, double vf_clip_coef, double vf_coef, double ent_coef, int T_seq, int A, int N, cudaStream_t stream) {
    launch_ppo_loss_forward<float>(loss_output, saved_for_backward, logits, values_pred, actions, old_logprobs, advantages, prio, values, returns, adv_mean, adv_std, clip_coef, vf_clip_coef, vf_coef, ent_coef, T_seq, A, N, stream);
}
void launch_ppo_loss_backward_float(float* grad_logits, float* grad_values_pred, const float* grad_loss, const float* logits, const int64_t* actions, const float* old_logprobs, const float* advantages, const float* prio, const float* values, const float* returns, const double* saved_for_backward, const float* adv_mean, const float* adv_std, double clip_coef, double vf_clip_coef, double vf_coef, double ent_coef, int T_seq, int A, int N, cudaStream_t stream) {
    launch_ppo_loss_backward<float>(grad_logits, grad_values_pred, grad_loss, logits, actions, old_logprobs, advantages, prio, values, returns, saved_for_backward, adv_mean, adv_std, clip_coef, vf_clip_coef, vf_coef, ent_coef, T_seq, A, N, stream);
}
void launch_sample_logits_float(double* actions, float* logprobs, float* value_out, const float* logits, const float* value, uint64_t seed, const int64_t* offset_ptr, int A, int B, int logits_stride, int value_stride, cudaStream_t stream) {
    launch_sample_logits<float>(actions, logprobs, value_out, logits, value, seed, offset_ptr, A, B, logits_stride, value_stride, stream);
}

// Non-templated wrappers for BFloat16
void launch_mingru_gate_inference_bf16(at::BFloat16* out, at::BFloat16* next_state, const at::BFloat16* combined, const at::BFloat16* state_in, int H, int B, cudaStream_t stream) {
    launch_mingru_gate_inference<at::BFloat16>(out, next_state, combined, state_in, H, B, stream);
}
void launch_log_coeffs_and_values_bf16(at::BFloat16* log_coeffs, at::BFloat16* log_values, const at::BFloat16* gate, const at::BFloat16* hidden, int N, cudaStream_t stream) {
    launch_log_coeffs_and_values<at::BFloat16>(log_coeffs, log_values, gate, hidden, N, stream);
}
void launch_log_coeffs_and_values_backward_bf16(at::BFloat16* grad_gate, at::BFloat16* grad_hidden, const at::BFloat16* grad_log_coeffs, const at::BFloat16* grad_log_values, const at::BFloat16* gate, const at::BFloat16* hidden, int N, cudaStream_t stream) {
    launch_log_coeffs_and_values_backward<at::BFloat16>(grad_gate, grad_hidden, grad_log_coeffs, grad_log_values, gate, hidden, N, stream);
}
void launch_rmsnorm_forward_bf16(at::BFloat16* out, float* inv_norm_buf, const at::BFloat16* x, const at::BFloat16* weight, double eps, int T_total, int H, int B, cudaStream_t stream) {
    launch_rmsnorm_forward<at::BFloat16>(out, inv_norm_buf, x, weight, eps, T_total, H, B, stream);
}
void launch_rmsnorm_backward_bf16(at::BFloat16* grad_x, at::BFloat16* grad_weight, const at::BFloat16* grad_out, const float* inv_norm_buf, const at::BFloat16* x_buf, const at::BFloat16* weight, double eps, int T_total, int H, int B, cudaStream_t stream) {
    launch_rmsnorm_backward<at::BFloat16>(grad_x, grad_weight, grad_out, inv_norm_buf, x_buf, weight, eps, T_total, H, B, stream);
}
void launch_fused_scan_forward_bf16(at::BFloat16* out, at::BFloat16* next_state, float* a_star, float* s_vals, float* log_values_buf, const at::BFloat16* combined, const at::BFloat16* state, int T_seq, int H, int B, cudaStream_t stream) {
    launch_fused_scan_forward<at::BFloat16>(out, next_state, a_star, s_vals, log_values_buf, combined, state, T_seq, H, B, stream);
}
void launch_fused_scan_backward_bf16(at::BFloat16* grad_combined, at::BFloat16* grad_state, const at::BFloat16* grad_out, const at::BFloat16* grad_next_state, const at::BFloat16* combined, const at::BFloat16* state, const float* a_star_buf, const float* s_buf, const float* log_values_buf, int T_seq, int H, int B, cudaStream_t stream) {
    launch_fused_scan_backward<at::BFloat16>(grad_combined, grad_state, grad_out, grad_next_state, combined, state, a_star_buf, s_buf, log_values_buf, T_seq, H, B, stream);
}
void launch_logcumsumexp_forward_bf16(at::BFloat16* out, double* s_buf, const at::BFloat16* x, int T_total, int H, int B, cudaStream_t stream) {
    launch_logcumsumexp_forward<at::BFloat16>(out, s_buf, x, T_total, H, B, stream);
}
void launch_logcumsumexp_backward_bf16(at::BFloat16* grad_x, const at::BFloat16* grad_out, const at::BFloat16* x, const double* s_buf, int T_total, int H, int B, cudaStream_t stream) {
    launch_logcumsumexp_backward<at::BFloat16>(grad_x, grad_out, x, s_buf, T_total, H, B, stream);
}
void launch_ppo_loss_forward_bf16(float* loss_output, double* saved_for_backward, const at::BFloat16* logits, const at::BFloat16* values_pred, const int64_t* actions, const at::BFloat16* old_logprobs, const at::BFloat16* advantages, const at::BFloat16* prio, const at::BFloat16* values, const at::BFloat16* returns, const float* adv_mean, const float* adv_std, double clip_coef, double vf_clip_coef, double vf_coef, double ent_coef, int T_seq, int A, int N, cudaStream_t stream) {
    launch_ppo_loss_forward<at::BFloat16>(loss_output, saved_for_backward, logits, values_pred, actions, old_logprobs, advantages, prio, values, returns, adv_mean, adv_std, clip_coef, vf_clip_coef, vf_coef, ent_coef, T_seq, A, N, stream);
}
void launch_ppo_loss_backward_bf16(at::BFloat16* grad_logits, at::BFloat16* grad_values_pred, const float* grad_loss, const at::BFloat16* logits, const int64_t* actions, const at::BFloat16* old_logprobs, const at::BFloat16* advantages, const at::BFloat16* prio, const at::BFloat16* values, const at::BFloat16* returns, const double* saved_for_backward, const float* adv_mean, const float* adv_std, double clip_coef, double vf_clip_coef, double vf_coef, double ent_coef, int T_seq, int A, int N, cudaStream_t stream) {
    launch_ppo_loss_backward<at::BFloat16>(grad_logits, grad_values_pred, grad_loss, logits, actions, old_logprobs, advantages, prio, values, returns, saved_for_backward, adv_mean, adv_std, clip_coef, vf_clip_coef, vf_coef, ent_coef, T_seq, A, N, stream);
}
void launch_sample_logits_bf16(double* actions, at::BFloat16* logprobs, at::BFloat16* value_out, const at::BFloat16* logits, const at::BFloat16* value, uint64_t seed, const int64_t* offset_ptr, int A, int B, int logits_stride, int value_stride, cudaStream_t stream) {
    launch_sample_logits<at::BFloat16>(actions, logprobs, value_out, logits, value, seed, offset_ptr, A, B, logits_stride, value_stride, stream);
}

#endif // PUFFERLIB_KERNELS_CU