932ac2bda8
Two bugs caught by the L40S smoke (train-qhgj6) that couldn't surface on
local RTX-3050 single-fold runs:
1. PER dtoh inside CUDA Graph capture (Fold 1 crash)
Failure: CUDA_ERROR_STREAM_CAPTURE_INVALIDATED at per_prefix_scan on
Fold 1 re-capture. Chain: fused_training parent graph captures →
memcpy_dtoh + cuStreamSynchronize in gpu_replay_buffer::update_priorities_gpu
(health<0.8 diversity path) poisons the stream → subsequent per_sample
kernel on the same stream sees an invalidated capture context.
The prior comment claimed "runs once per epoch, DtoH cost acceptable"
— wrong, it runs every priority update when health<0.8 (common during
Fold handoff when health_cache is re-seeded low). Any dtoh inside
capture invalidates regardless of latency.
Proper fix (no shortcut):
* New kernel actions_sum_scale_reduce_u32 — single-block deterministic
tree reduction over sample_actions (u32) → writes (sum*1000)/n as i32
to a device-accessible slot. No atomics (consistent with the 1/N
determinism policy from commit c82386500).
* mean_action_scaled storage is pinned + device-mapped (cuMemAllocHost
+ cuMemHostGetDevicePointer — same pattern as rng_step_dev_ptr and
size_dev_ptr elsewhere in the file). Zero-copy between host and
device, graph-safe, no explicit free needed (process-exit cleanup,
matches existing pattern).
* pow_alpha_diverse_f32 now takes const int* mean_action_scaled_ptr
and does a plain global load — NOT __ldg. The read-only cache used
by __ldg is not guaranteed coherent with device-mapped host memory;
multi-trial smoke regression caught it (median q_gap collapsed
from 2.0 → 0.15 with __ldg, recovered to 2.8 with plain load).
Verified: multi-trial smoke 5/5 pass, median_q_gap=2.80 (beats 2.00
baseline), Best Sharpe peaks 19-30 per trial. No stream capture
invalidation.
2. evaluate step CLI drift in Argo template
evaluate_baseline's Args struct uses --models-dir and --output (single
file path). Template was passing --checkpoint-dir and --output-dir,
causing clap to reject the invocation. Fixed argument names + added
mkdir for the eval subdir + updated the comment to pin the source of
truth for future drift catches.
Both fixes are graph-capture-clean and match the "wire properly or delete"
discipline. No masking, no feature flags, no dead params.
593 lines
20 KiB
Plaintext
593 lines
20 KiB
Plaintext
// CUDA kernels for GPU-resident replay buffer (pure cudarc, no Candle).
|
||
//
|
||
// Kernels:
|
||
// 1. scatter_insert_f32 — ring buffer insert for f32 arrays (rewards, dones, priorities)
|
||
// 2. scatter_insert_u32 — ring buffer insert for u32 arrays (actions)
|
||
// 3. scatter_insert — ring buffer insert for state matrices
|
||
// 4. gather_f32 — index_select for f32 arrays
|
||
// 5. gather_u32 — index_select for u32 arrays
|
||
// 6. gather_rows — index_select rows from [cap, dim] matrix
|
||
// 7. pow_alpha_f32 — prios[i] = pow(abs(input[i]), alpha)
|
||
// 8. is_weights_f32 — importance sampling weight computation
|
||
// 9. priority_update_f32 — |td|^alpha + eps scatter-write + atomicMax
|
||
// 10. fill_f32 — fill array with scalar value
|
||
// 11. rand_scale_f32 — scale uniform [0,1) by total_sum
|
||
|
||
// ── 1. Scatter insert for f32 (1D) ─────────────────────────────────────────
|
||
|
||
extern "C" __global__
|
||
void scatter_insert_f32(
|
||
float* __restrict__ dst, // [capacity]
|
||
const float* __restrict__ src, // [batch_size]
|
||
int start_idx, // write cursor position
|
||
int capacity, // ring buffer capacity
|
||
int batch_size // number of elements to write
|
||
) {
|
||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||
if (i >= batch_size) return;
|
||
int dst_idx = (start_idx + i) % capacity;
|
||
dst[dst_idx] = src[i];
|
||
}
|
||
|
||
// ── 2. Scatter insert for u32 (1D) ─────────────────────────────────────────
|
||
|
||
extern "C" __global__
|
||
void scatter_insert_u32(
|
||
unsigned int* __restrict__ dst,
|
||
const unsigned int* __restrict__ src,
|
||
int start_idx,
|
||
int capacity,
|
||
int batch_size
|
||
) {
|
||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||
if (i >= batch_size) return;
|
||
int dst_idx = (start_idx + i) % capacity;
|
||
dst[dst_idx] = src[i];
|
||
}
|
||
|
||
// ── 3. Scatter insert for f32 rows [batch, dim] → ring buffer [cap, dim] ──
|
||
// 2D-aware version of scatter_insert_f32 for state matrices.
|
||
// Each thread writes one element: dst[(start + row) % cap, col] = src[row, col].
|
||
|
||
extern "C" __global__
|
||
void scatter_insert_f32_rows(
|
||
float* __restrict__ dst, // [capacity * state_dim]
|
||
const float* __restrict__ src, // [batch_size * state_dim]
|
||
int start_idx,
|
||
int capacity,
|
||
int state_dim,
|
||
int batch_size
|
||
) {
|
||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||
int total = batch_size * state_dim;
|
||
if (tid >= total) return;
|
||
int row = tid / state_dim;
|
||
int col = tid % state_dim;
|
||
int dst_row = (start_idx + row) % capacity;
|
||
dst[dst_row * state_dim + col] = src[row * state_dim + col];
|
||
}
|
||
|
||
// ── 4. Gather f32 by indices ────────────────────────────────────────────────
|
||
|
||
extern "C" __global__
|
||
void gather_f32(
|
||
float* __restrict__ out, // [batch_size]
|
||
const float* __restrict__ src, // [capacity]
|
||
const long long* __restrict__ indices, // [batch_size]
|
||
int batch_size
|
||
) {
|
||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||
if (i >= batch_size) return;
|
||
int idx = (int)indices[i];
|
||
out[i] = src[idx];
|
||
}
|
||
|
||
// ── 5. Gather u32 by indices ────────────────────────────────────────────────
|
||
|
||
extern "C" __global__
|
||
void gather_u32(
|
||
unsigned int* __restrict__ out,
|
||
const unsigned int* __restrict__ src,
|
||
const long long* __restrict__ indices,
|
||
int batch_size,
|
||
int capacity
|
||
) {
|
||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||
if (i >= batch_size) return;
|
||
int idx = (int)indices[i];
|
||
if (idx < 0 || idx >= capacity) return;
|
||
out[i] = src[idx];
|
||
}
|
||
|
||
// ── 6. Gather rows [cap, dim] → [batch, dim] by indices ───────────────
|
||
|
||
extern "C" __global__
|
||
void gather_rows(
|
||
unsigned short* __restrict__ out, // [batch_size * state_dim]
|
||
const unsigned short* __restrict__ src, // [capacity * state_dim]
|
||
const long long* __restrict__ indices, // [batch_size]
|
||
int state_dim,
|
||
int batch_size
|
||
) {
|
||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||
int total = batch_size * state_dim;
|
||
if (tid >= total) return;
|
||
int row = tid / state_dim;
|
||
int col = tid % state_dim;
|
||
int src_row = (int)indices[row];
|
||
out[row * state_dim + col] = src[src_row * state_dim + col];
|
||
}
|
||
|
||
// ── 6a. Gather f32 rows [cap, dim] → [batch, dim] by indices (#30) ─────────
|
||
|
||
extern "C" __global__
|
||
void gather_f32_rows(
|
||
float* __restrict__ out, // [batch_size * state_dim]
|
||
const float* __restrict__ src, // [capacity * state_dim]
|
||
const long long* __restrict__ indices, // [batch_size]
|
||
int state_dim,
|
||
int batch_size
|
||
) {
|
||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||
int total = batch_size * state_dim;
|
||
if (tid >= total) return;
|
||
int row = tid / state_dim;
|
||
int col = tid % state_dim;
|
||
int src_row = (int)indices[row];
|
||
out[row * state_dim + col] = src[src_row * state_dim + col];
|
||
}
|
||
|
||
// ── 6b. Gather scalars (1D) by indices ────────────────────────────────
|
||
|
||
extern "C" __global__
|
||
void gather(
|
||
unsigned short* __restrict__ out, // [batch_size]
|
||
const unsigned short* __restrict__ src, // [capacity]
|
||
const long long* __restrict__ indices, // [batch_size]
|
||
int batch_size,
|
||
int capacity // bounds check for index validity
|
||
) {
|
||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||
if (i >= batch_size) return;
|
||
int idx = (int)indices[i];
|
||
/* Bounds check: corrupt PER segment tree can produce OOB indices.
|
||
* Clamp to [0, capacity-1] instead of reading garbage (NaN bits). */
|
||
if (idx < 0 || idx >= capacity) return;
|
||
out[i] = src[idx];
|
||
}
|
||
|
||
// ── 7. prios_alpha: out[i] = pow(src[i], alpha) ────────────────────────────
|
||
|
||
extern "C" __global__
|
||
void pow_alpha_f32(
|
||
float* __restrict__ out,
|
||
const float* __restrict__ src,
|
||
float alpha,
|
||
int n
|
||
) {
|
||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||
if (i >= n) return;
|
||
out[i] = powf(src[i], alpha);
|
||
}
|
||
|
||
// ── 8. IS weights: w_i = (N * sampled_prio[i] / total_sum)^(-beta) ─────────
|
||
// Then normalize by max weight.
|
||
// total_sum_buf is a GPU-resident scalar pointer (zero CPU readback).
|
||
|
||
extern "C" __global__
|
||
void is_weights_f32(
|
||
float* __restrict__ weights, // [batch_size] output (in-place)
|
||
const float* __restrict__ sampled_prios, // [batch_size]
|
||
const float* __restrict__ total_sum_buf, // [1] GPU-resident scalar
|
||
float neg_beta,
|
||
const int* __restrict__ n_buffer_ptr, // pinned device-mapped — current buffer size at replay
|
||
int batch_size
|
||
) {
|
||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||
if (i >= batch_size) return;
|
||
int n_buffer = *n_buffer_ptr;
|
||
float ts = total_sum_buf[0];
|
||
float prob = fmaxf((sampled_prios[i] * (float)n_buffer) / fmaxf(ts, 1e-8f), 1e-12f);
|
||
float w = powf(prob, neg_beta);
|
||
weights[i] = fminf(w, 60000.0f);
|
||
}
|
||
|
||
// Normalize weights by max (two-pass: first find max, then divide)
|
||
extern "C" __global__
|
||
void normalize_weights_f32(
|
||
float* __restrict__ weights,
|
||
const float* __restrict__ max_weight, // [1] scalar
|
||
int batch_size
|
||
) {
|
||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||
if (i >= batch_size) return;
|
||
float mw = max_weight[0];
|
||
if (mw < 1e-8f) mw = 1e-8f;
|
||
weights[i] /= mw;
|
||
}
|
||
|
||
// ── 9. Priority update: |td|^alpha + eps, scatter, atomicMax ────────────────
|
||
|
||
extern "C" __global__
|
||
void priority_update_f32(
|
||
const float* __restrict__ td_errors,
|
||
const unsigned int* __restrict__ indices,
|
||
float* __restrict__ priorities,
|
||
float* __restrict__ batch_max,
|
||
float alpha,
|
||
float epsilon,
|
||
int batch_size
|
||
) {
|
||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||
if (i >= batch_size) return;
|
||
|
||
float td = fabsf(td_errors[i]);
|
||
float new_prio = powf(td, alpha) + epsilon;
|
||
|
||
// Clamp to [epsilon, 1e6]
|
||
if (new_prio < epsilon) new_prio = epsilon;
|
||
if (new_prio > 1e6f) new_prio = 1e6f;
|
||
|
||
unsigned int idx = indices[i];
|
||
priorities[idx] = new_prio;
|
||
|
||
// atomicMax for batch max: IEEE 754 positive floats preserve int ordering
|
||
int ival = __float_as_int(new_prio);
|
||
atomicMax((int*)batch_max, ival);
|
||
}
|
||
|
||
// ── 10. Fill f32 array with scalar ──────────────────────────────────────────
|
||
|
||
extern "C" __global__
|
||
void fill_f32(
|
||
float* __restrict__ out,
|
||
float value,
|
||
int n
|
||
) {
|
||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||
if (i >= n) return;
|
||
out[i] = value;
|
||
}
|
||
|
||
// ── 11. Scale uniform [0,1) by total_sum scalar ────────────────────────────
|
||
|
||
extern "C" __global__
|
||
void rand_scale_f32(
|
||
float* __restrict__ targets, // [batch_size] in-place: targets[i] *= total_sum
|
||
const float* __restrict__ total_sum, // [1]
|
||
int batch_size
|
||
) {
|
||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||
if (i >= batch_size) return;
|
||
targets[i] *= total_sum[0];
|
||
}
|
||
|
||
// ── 12. Reduce max f32 ─────────────────────────────────────────────────────
|
||
// Single-block parallel reduction to find max of array.
|
||
|
||
extern "C" __global__
|
||
void reduce_max_f32(
|
||
const float* __restrict__ input,
|
||
float* __restrict__ output, // [1]
|
||
int n
|
||
) {
|
||
extern __shared__ float sdata[];
|
||
|
||
int tid = threadIdx.x;
|
||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||
|
||
sdata[tid] = (i < n) ? input[i] : -1e30f;
|
||
__syncthreads();
|
||
|
||
for (int s = blockDim.x / 2; s > 0; s >>= 1) {
|
||
if (tid < s && sdata[tid] < sdata[tid + s]) {
|
||
sdata[tid] = sdata[tid + s];
|
||
}
|
||
__syncthreads();
|
||
}
|
||
|
||
if (tid == 0) {
|
||
// atomicMax across blocks (for multi-block launches)
|
||
int ival = __float_as_int(sdata[0]);
|
||
atomicMax((int*)output, ival);
|
||
}
|
||
}
|
||
|
||
// ── 13. Convert i64 indices to u32 ──────────────────────────────────────────
|
||
|
||
extern "C" __global__
|
||
void i64_to_u32(
|
||
unsigned int* __restrict__ out,
|
||
const long long* __restrict__ input,
|
||
int n
|
||
) {
|
||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||
if (i >= n) return;
|
||
out[i] = (unsigned int)input[i];
|
||
}
|
||
|
||
// ── 14. Cast F32 states to U16 for ring buffer storage ─────────────────────
|
||
|
||
extern "C" __global__
|
||
void f32_to_u16_cast(
|
||
unsigned short* __restrict__ out,
|
||
const float* __restrict__ input,
|
||
int n
|
||
) {
|
||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||
if (i >= n) return;
|
||
/* Reinterpret f32 bit pattern to u16. */
|
||
float val = (input[i]);
|
||
out[i] = *(unsigned short*)&val;
|
||
}
|
||
|
||
// ── 14b. Fill f32 array from a GPU scalar pointer ───────────────────────────
|
||
|
||
extern "C" __global__
|
||
void fill_from_gpu_f32(
|
||
float* __restrict__ out,
|
||
const float* __restrict__ value_ptr, // [1] scalar on GPU
|
||
int n
|
||
) {
|
||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||
if (i >= n) return;
|
||
out[i] = value_ptr[0];
|
||
}
|
||
|
||
// ── 15. Cast U16 to F32 on GPU ───────────────────────────────────────
|
||
// Inverse of f32_to_u16_cast. Converts u16 storage bits to f32.
|
||
|
||
extern "C" __global__
|
||
void u16_to_f32_cast(
|
||
float* __restrict__ out,
|
||
const unsigned short* __restrict__ input,
|
||
int n
|
||
) {
|
||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||
if (i >= n) return;
|
||
// Upper 16 bits of F32, so shift left by 16
|
||
unsigned int bits = ((unsigned int)input[i]) << 16;
|
||
out[i] = __uint_as_float(bits);
|
||
}
|
||
|
||
// ── 16. Cast u32 to F32 on GPU ──────────────────────────────────────────────
|
||
|
||
extern "C" __global__
|
||
void u32_to_f32_cast(
|
||
float* __restrict__ out,
|
||
const unsigned int* __restrict__ input,
|
||
int n
|
||
) {
|
||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||
if (i >= n) return;
|
||
out[i] = (float)input[i];
|
||
}
|
||
|
||
// ── 17. GPU max of two single-element f32 scalars ───────────────────────────
|
||
// out[0] = max(a[0], b[0])
|
||
|
||
extern "C" __global__
|
||
void max_of_two_f32(
|
||
float* __restrict__ out,
|
||
const float* __restrict__ a,
|
||
const float* __restrict__ b
|
||
) {
|
||
float va = a[0];
|
||
float vb = b[0];
|
||
out[0] = (va > vb) ? va : vb;
|
||
}
|
||
|
||
// ── 18a. Gather f32 rows with zero-padding to dst_stride ───────────────────
|
||
// Replaces 2-step gather + pad_states: gathers rows from [cap, src_dim]
|
||
// into [batch, dst_stride], zero-filling columns [src_dim, dst_stride).
|
||
|
||
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;
|
||
}
|
||
|
||
// ── 18b. Gather scalar f32 by i64 indices ─────────────────────────────────
|
||
// Direct gather for rewards, dones, IS-weights into trainer buffers.
|
||
|
||
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]];
|
||
}
|
||
|
||
// ── 18c. Gather scalar i32 by i64 indices ─────────────────────────────────
|
||
// Direct gather for actions (u32/i32) into trainer buffers.
|
||
|
||
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]];
|
||
}
|
||
|
||
// ── 18. NaN/Inf check kernel ────────────────────────────────────────────────
|
||
// out[i] = (isnan(v) || isinf(v)) ? 1.0f : 0.0f
|
||
|
||
extern "C" __global__
|
||
void nan_inf_check_f32(
|
||
float* __restrict__ out,
|
||
const float* __restrict__ input,
|
||
int n
|
||
) {
|
||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||
if (i >= n) return;
|
||
float v = input[i];
|
||
out[i] = (isnan(v) || isinf(v)) ? 1.0f : 0.0f;
|
||
}
|
||
|
||
// ── 19. Dead neuron (near-zero) check kernel ────────────────────────────────
|
||
// out[i] = (fabsf(v) < threshold) ? 1.0f : 0.0f
|
||
|
||
extern "C" __global__
|
||
void dead_neuron_check_f32(
|
||
float* __restrict__ out,
|
||
const float* __restrict__ input,
|
||
float threshold,
|
||
int n
|
||
) {
|
||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||
if (i >= n) return;
|
||
out[i] = (fabsf(input[i]) < threshold) ? 1.0f : 0.0f;
|
||
}
|
||
|
||
// ── 20. Cast f32-encoded indices to u32 ─────────────────────────────────────
|
||
// out[i] = (unsigned int)input[i]
|
||
|
||
extern "C" __global__
|
||
void f32_idx_to_u32(
|
||
unsigned int* __restrict__ out,
|
||
const float* __restrict__ input,
|
||
int n
|
||
) {
|
||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||
if (i >= n) return;
|
||
out[i] = (unsigned int)input[i];
|
||
}
|
||
|
||
/* Graph-safe GPU reduction: scaled-mean of u32 actions → i32[1].
|
||
*
|
||
* Writes (sum(actions) * 1000) / n into mean_scaled_out[0]. Launched as a
|
||
* single 256-thread block; deterministic tree reduction in shared memory.
|
||
* Replaces the prior dtoh-then-sum-on-host path, which was not graph-safe:
|
||
* memcpy_dtoh inside a CUDA Graph capture region invalidates the stream
|
||
* (L40S smoke surfaced this on Fold 1 re-capture). Keep n small — called
|
||
* with batch-sized windows (≤ max_batch_size), one block handles the lot.
|
||
*/
|
||
extern "C" __global__ void actions_sum_scale_reduce_u32(
|
||
const unsigned int* __restrict__ actions,
|
||
int* __restrict__ mean_scaled_out, /* i32[1]: (sum * 1000) / n */
|
||
int n
|
||
) {
|
||
extern __shared__ long long sdata_ll[];
|
||
int tid = threadIdx.x;
|
||
|
||
long long local = 0;
|
||
for (int i = tid; i < n; i += blockDim.x) {
|
||
local += (long long)actions[i];
|
||
}
|
||
sdata_ll[tid] = local;
|
||
__syncthreads();
|
||
|
||
/* Deterministic tree reduction — no atomics, no warp shuffles with
|
||
* data-dependent lane ordering. */
|
||
for (int offset = blockDim.x / 2; offset > 0; offset >>= 1) {
|
||
if (tid < offset) {
|
||
sdata_ll[tid] += sdata_ll[tid + offset];
|
||
}
|
||
__syncthreads();
|
||
}
|
||
|
||
if (tid == 0) {
|
||
long long denom = n > 0 ? (long long)n : 1LL;
|
||
long long mean_scaled = (sdata_ll[0] * 1000LL) / denom;
|
||
mean_scaled_out[0] = (int)mean_scaled;
|
||
}
|
||
}
|
||
|
||
/* C1/P1: Health-weighted PER priority — boosts priorities of diverse-action
|
||
* experiences during collapse (health < 0.8).
|
||
*
|
||
* new_prio = clamp((|td_error[i]| + epsilon)^alpha
|
||
* × (1 + 2×(1−health) × |action[i] − mean_action|),
|
||
* epsilon, 1e6)
|
||
* priorities[indices[i]] = new_prio
|
||
* priorities_pa[indices[i]] = new_prio^alpha
|
||
* atomicMax(batch_max, new_prio) (IEEE 754 int trick)
|
||
*
|
||
* `mean_action_scaled_ptr` points to a device-side i32 set by the companion
|
||
* kernel `actions_sum_scale_reduce_u32` (same batch, scaled mean × 1000).
|
||
* Device-pointer passthrough keeps the update graph-safe — no host readback.
|
||
*
|
||
* Drop-in replacement for per_update_pa when health < 0.8.
|
||
*/
|
||
extern "C" __global__ void pow_alpha_diverse_f32(
|
||
float* __restrict__ priorities_pa,
|
||
float* __restrict__ priorities,
|
||
float* __restrict__ batch_max,
|
||
const unsigned int* __restrict__ indices,
|
||
const float* __restrict__ td_errors,
|
||
const int* __restrict__ actions,
|
||
const int* __restrict__ mean_action_scaled_ptr, /* device i32[1] */
|
||
float alpha,
|
||
float epsilon,
|
||
float health,
|
||
int capacity,
|
||
int batch_size
|
||
) {
|
||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||
if (i >= batch_size) return;
|
||
|
||
unsigned int idx = indices[i];
|
||
if (idx >= (unsigned int)capacity) return;
|
||
|
||
float base = powf(fabsf(td_errors[i]) + epsilon, alpha);
|
||
|
||
/* Plain load (not __ldg): the backing buffer is device-mapped pinned host
|
||
* memory (cuMemHostAlloc + cuMemHostGetDevicePointer). __ldg routes through
|
||
* the read-only cache which is not guaranteed coherent for mapped host
|
||
* memory — stale reads manifested as a 13× drop in final q_gap on the
|
||
* multi-trial smoke. Plain load respects stream serialization between
|
||
* actions_sum_scale_reduce_u32 (writer) and this kernel (reader). */
|
||
float mean_action = (float)(mean_action_scaled_ptr[0]) * 0.001f;
|
||
float action_diff = fabsf((float)actions[i] - mean_action);
|
||
float diversity_mult = 1.0f + 2.0f * (1.0f - health) * action_diff;
|
||
|
||
float new_prio = base * diversity_mult;
|
||
if (new_prio < epsilon) new_prio = epsilon;
|
||
if (new_prio > 1e6f) new_prio = 1e6f;
|
||
|
||
priorities[idx] = new_prio;
|
||
|
||
int ival = __float_as_int(new_prio);
|
||
atomicMax((int*)batch_max, ival);
|
||
|
||
float pa = powf(new_prio, alpha);
|
||
priorities_pa[idx] = pa;
|
||
}
|
||
|
||
// ── 21. Regime weight combination kernel ────────────────────────────────────
|
||
// out[i] = trending[i] * t_scale + ranging[i] * r_scale + volatile[i] * v_scale
|
||
|
||
extern "C" __global__
|
||
void regime_weight_combine_f32(
|
||
float* __restrict__ out,
|
||
const float* __restrict__ trending,
|
||
const float* __restrict__ ranging,
|
||
const float* __restrict__ volatile_mask,
|
||
float t_scale,
|
||
float r_scale,
|
||
float v_scale,
|
||
int n
|
||
) {
|
||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||
if (i >= n) return;
|
||
out[i] = trending[i] * t_scale + ranging[i] * r_scale + volatile_mask[i] * v_scale;
|
||
}
|