feat: graph_utility_kernels.cu — gather_padded + increment_step_counters for true single-graph

Adds 4 GPU-native kernels replacing host-side operations on the training
hot path: gather_f32_rows_padded (row gather + 128-byte zero-pad),
gather_f32_scalar, gather_i32_scalar, and increment_step_counters (atomic
counter bumps + cosine-annealed tau — zero CPU sync per step).
Wired into build.rs (nvcc cubin compile) and gpu_dqn_trainer.rs
(GRAPH_UTILITY_CUBIN include_bytes!).

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

crates/ml/build.rs

@@ -81,6 +81,7 @@ fn main() {
         "backward_kernels.cu",
         "iqn_cvar_kernel.cu",
         "mamba2_temporal_kernel.cu",
+        "graph_utility_kernels.cu",
     ];
 
     // ALL kernels get common header (BF16 types + wrappers)

crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs

@@ -71,6 +71,7 @@ static Q_STATS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/q_stats_
 static ENSEMBLE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/ensemble_kernels.cubin"));
 static CQL_GRAD_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/cql_grad_kernel.cubin"));
 static MAMBA2_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/mamba2_temporal_kernel.cubin"));
+pub(crate) static GRAPH_UTILITY_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/graph_utility_kernels.cubin"));
 
 /// Mamba2 temporal scan configuration.
 const MAMBA2_HISTORY_K: usize = 8;    // Rolling history length

crates/ml/src/cuda_pipeline/graph_utility_kernels.cu

@@ -0,0 +1,146 @@
+/**
+ * Graph utility kernels — GPU-native gather + step counter increment.
+ *
+ * These kernels replace host-side operations on the training hot path,
+ * enabling a true single-graph architecture with zero CPU involvement
+ * per step (single cuGraphExecLaunch).
+ *
+ * Kernels:
+ *   gather_f32_rows_padded   — gather rows from replay buffer + zero-pad
+ *                              to 128-byte-aligned dst_stride (replaces
+ *                              2-step gather + pad_states host pipeline)
+ *   gather_f32_scalar        — gather scalar f32 values (rewards, dones,
+ *                              IS-weights) from replay buffer indices
+ *   gather_i32_scalar        — gather scalar i32 values (actions) from
+ *                              replay buffer indices
+ *   increment_step_counters  — single-thread kernel that increments all
+ *                              per-step Adam/scheduler counters on GPU
+ *                              and computes cosine-annealed tau
+ *
+ * No #include of common_device_functions.cuh is needed here — the build
+ * function in build.rs prepends it via string concatenation before
+ * passing the combined source to nvcc.
+ */
+
+/* ══════════════════════════════════════════════════════════════════════
+ * gather_f32_rows_padded
+ *
+ * Gathers rows from a flat row-major src buffer into dst, zero-padding
+ * columns [src_dim, dst_stride) so the output is 128-byte aligned per row.
+ *
+ * Launch config: grid=(ceil(batch_size * dst_stride / 256)), block=(256)
+ *
+ * Parameters:
+ *   dst        — output buffer [batch_size * dst_stride]
+ *   src        — source buffer [capacity * src_dim]
+ *   indices    — i64 indices into src rows [batch_size]
+ *   src_dim    — number of valid f32 elements per src row (unpadded)
+ *   dst_stride — padded row width in f32 elements (dst_dim >= src_dim,
+ *                aligned to 128 bytes = 32 f32 elements)
+ *   batch_size — number of rows to gather
+ * ══════════════════════════════════════════════════════════════════════ */
+extern "C" __global__ void gather_f32_rows_padded(
+    float* __restrict__ dst,
+    const float* __restrict__ src,
+    const long long* __restrict__ indices,
+    int src_dim, int dst_stride, int batch_size)
+{
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+    int total = batch_size * dst_stride;
+    if (tid >= total) return;
+    int row = tid / dst_stride;
+    int col = tid % dst_stride;
+    long long src_row = indices[row];
+    dst[tid] = (col < src_dim) ? src[src_row * src_dim + col] : 0.0f;
+}
+
+/* ══════════════════════════════════════════════════════════════════════
+ * gather_f32_scalar
+ *
+ * Gathers scalar f32 values (rewards, dones, IS-weights) from the replay
+ * buffer into a flat output array using the sampled indices.
+ *
+ * Launch config: grid=(ceil(batch_size / 256)), block=(256)
+ *
+ * Parameters:
+ *   dst        — output buffer [batch_size]
+ *   src        — source buffer [capacity] (replay buffer scalar field)
+ *   indices    — i64 indices [batch_size]
+ *   batch_size — number of elements to gather
+ * ══════════════════════════════════════════════════════════════════════ */
+extern "C" __global__ void gather_f32_scalar(
+    float* __restrict__ dst,
+    const float* __restrict__ src,
+    const long long* __restrict__ indices,
+    int batch_size)
+{
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+    if (tid >= batch_size) return;
+    dst[tid] = src[indices[tid]];
+}
+
+/* ══════════════════════════════════════════════════════════════════════
+ * gather_i32_scalar
+ *
+ * Gathers scalar i32 values (packed actions) from the replay buffer into
+ * a flat output array using the sampled indices.
+ *
+ * Launch config: grid=(ceil(batch_size / 256)), block=(256)
+ *
+ * Parameters:
+ *   dst        — output buffer [batch_size]
+ *   src        — source buffer [capacity] (replay buffer action field)
+ *   indices    — i64 indices [batch_size]
+ *   batch_size — number of elements to gather
+ * ══════════════════════════════════════════════════════════════════════ */
+extern "C" __global__ void gather_i32_scalar(
+    int* __restrict__ dst,
+    const int* __restrict__ src,
+    const long long* __restrict__ indices,
+    int batch_size)
+{
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+    if (tid >= batch_size) return;
+    dst[tid] = src[indices[tid]];
+}
+
+/* ══════════════════════════════════════════════════════════════════════
+ * increment_step_counters
+ *
+ * Single-thread kernel that atomically increments all per-step counters
+ * and recomputes the cosine-annealed IQN tau value — all on GPU, zero
+ * CPU synchronisation required.
+ *
+ * All counter pointers must be device-accessible (device alloc or pinned
+ * host-mapped memory).  tau_out is written on every call.
+ *
+ * Launch config: grid=(1), block=(1)
+ *
+ * Parameters:
+ *   adam_t, sel_t, denoise_t, iql_t, iql_low_t, iqn_t, attn_t, rng_step
+ *              — per-scheduler step counters (int*, device-mapped)
+ *   tau_out    — cosine-annealed tau written here (float*, device-mapped)
+ *   tau_init   — initial tau value (start of annealing)
+ *   tau_final  — final tau value (end of annealing)
+ *   anneal_steps — total number of steps over which to anneal
+ * ══════════════════════════════════════════════════════════════════════ */
+extern "C" __global__ void increment_step_counters(
+    int* adam_t, int* sel_t, int* denoise_t, int* iql_t,
+    int* iql_low_t, int* iqn_t, int* attn_t, int* rng_step,
+    float* tau_out, float tau_init, float tau_final, int anneal_steps)
+{
+    /* adam_t is the master step — its post-increment value drives the
+     * cosine schedule so that step 1 produces the first non-initial tau. */
+    int step = atomicAdd(adam_t, 1) + 1;
+    atomicAdd(sel_t,      1);
+    atomicAdd(denoise_t,  1);
+    atomicAdd(iql_t,      1);
+    atomicAdd(iql_low_t,  1);
+    atomicAdd(iqn_t,      1);
+    atomicAdd(attn_t,     1);
+    atomicAdd(rng_step,   1);
+
+    float progress = fminf((float)step / (float)anneal_steps, 1.0f);
+    float cosine   = 0.5f * (1.0f + cosf(progress * 3.14159265f));
+    *tau_out = tau_final + (tau_init - tau_final) * cosine;
+}