feat(gpu): warp-cooperative DQN experience kernel for H100 (sm_90+)
Add warp-cooperative CUDA kernel for DQN experience collection that exploits H100's improved warp shuffle throughput. On sm_90+, 32 threads (1 warp) cooperate on each dot product via __shfl_xor_sync butterfly reduction, cutting per-thread register pressure from ~5.5KB to ~200B and enabling full occupancy on 132 SMs. Key additions: - 6 warp-cooperative device functions in common_device_functions.cuh: distributed/broadcast matvec (clean + NoisyNet), warp RMSNorm, warp_reduce_sum_all - 4 TILE_LAYER_WARP_* macros using __syncwarp() for warp-level sync - 3 warp forward passes: standard dueling, NoisyNet, C51 distributional (distributed heavy layers + broadcast atom layers for softmax) - Full warp kernel: dqn_full_experience_kernel_warp with lane-0 simulation, strided state scatter, warp-shuffle curiosity gather - Host-side compute capability detection (sm_major >= 9) with automatic fallback to standard 256-thread kernel on older GPUs - cuCtxSetLimit(STACK_SIZE, 16KB) for standard kernel safety 874/874 tests pass, 0 clippy warnings. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -258,6 +258,245 @@ __device__ void noisy_matvec_leaky_relu_shmem(
|
||||
}
|
||||
}
|
||||
|
||||
/* ------------------------------------------------------------------ */
|
||||
/* Warp-Cooperative Matrix-Vector Multiply */
|
||||
/* ------------------------------------------------------------------ */
|
||||
|
||||
/**
|
||||
* Warp-cooperative reduce: sum across 32 lanes using shuffle-down.
|
||||
* Result is valid in ALL lanes (full butterfly reduction).
|
||||
*/
|
||||
__device__ float warp_reduce_sum_all(float val) {
|
||||
for (int offset = 16; offset > 0; offset >>= 1)
|
||||
val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
|
||||
return val;
|
||||
}
|
||||
|
||||
/**
|
||||
* Warp-cooperative matrix-vector multiply with LeakyReLU, shared-memory tiled.
|
||||
*
|
||||
* Unlike the per-thread version, a full warp (32 lanes) cooperates on each
|
||||
* dot product. Input and output vectors are distributed across lanes in
|
||||
* strided layout: lane k holds elements k, k+32, k+64, ...
|
||||
*
|
||||
* This reduces per-thread local memory from in_dim floats (~2KB for dim=512)
|
||||
* to ceil(in_dim/32) floats (~64B), eliminating register spills on H100.
|
||||
*
|
||||
* @param shmem_W Shared memory weight tile [tile_rows × in_dim]
|
||||
* @param shmem_b Shared memory bias tile [tile_rows]
|
||||
* @param input_dist Per-lane distributed input [ceil(in_dim/32)]
|
||||
* @param output_dist Per-lane distributed output [ceil(out_dim/32)] (written)
|
||||
* @param in_dim Full input dimension
|
||||
* @param out_dim Full output dimension (only used for bounds context)
|
||||
* @param tile_offset Global output offset for this tile
|
||||
* @param tile_rows Number of output rows in this tile (≤ SHMEM_TILE_ROWS)
|
||||
* @param activate 1 = apply LeakyReLU, 0 = linear
|
||||
* @param lane_id threadIdx.x % 32
|
||||
*/
|
||||
__device__ void warp_matvec_leaky_relu_shmem(
|
||||
const float* __restrict__ shmem_W,
|
||||
const float* __restrict__ shmem_b,
|
||||
const float* input_dist,
|
||||
float* output_dist,
|
||||
int in_dim,
|
||||
int out_dim,
|
||||
int tile_offset,
|
||||
int tile_rows,
|
||||
int activate,
|
||||
int lane_id
|
||||
) {
|
||||
(void)out_dim;
|
||||
for (int j = 0; j < tile_rows; j++) {
|
||||
const float* row = shmem_W + j * in_dim;
|
||||
/* Each lane computes partial dot product over its elements */
|
||||
float partial = 0.0f;
|
||||
for (int i = lane_id; i < in_dim; i += 32)
|
||||
partial += row[i] * input_dist[i / 32];
|
||||
/* Add bias (only once, via lane 0) */
|
||||
if (lane_id == 0)
|
||||
partial += shmem_b[j];
|
||||
/* Full warp reduction — result in all lanes */
|
||||
float sum = warp_reduce_sum_all(partial);
|
||||
/* Activation */
|
||||
if (activate) sum = leaky_relu(sum);
|
||||
/* Scatter to owning lane in strided layout */
|
||||
int gj = tile_offset + j;
|
||||
if (lane_id == (gj & 31))
|
||||
output_dist[gj >> 5] = sum;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Warp-cooperative noisy matrix-vector multiply with factorized Gaussian noise.
|
||||
*
|
||||
* Same as warp_matvec_leaky_relu_shmem but applies NoisyNet perturbation.
|
||||
* Noise vectors are generated on-the-fly per lane (no large eps_in array needed).
|
||||
*
|
||||
* @param sigma_init NoisyNet sigma_0 parameter
|
||||
* @param rng Per-thread RNG state pointer
|
||||
*/
|
||||
__device__ void warp_noisy_matvec_leaky_relu_shmem(
|
||||
const float* __restrict__ shmem_W,
|
||||
const float* __restrict__ shmem_b,
|
||||
const float* input_dist,
|
||||
float* output_dist,
|
||||
int in_dim,
|
||||
int out_dim,
|
||||
int tile_offset,
|
||||
int tile_rows,
|
||||
int activate,
|
||||
float sigma_init,
|
||||
unsigned int* rng,
|
||||
int lane_id
|
||||
) {
|
||||
(void)out_dim;
|
||||
float sigma_scale = sigma_init / sqrtf((float)in_dim);
|
||||
|
||||
/* Each lane generates its portion of input noise (strided) */
|
||||
/* With in_dim=512 and 32 lanes, each lane generates 16 noise values */
|
||||
int local_count = (in_dim + 31) / 32;
|
||||
|
||||
for (int j = 0; j < tile_rows; j++) {
|
||||
/* Generate output noise for this row — same value needed by all lanes */
|
||||
/* Use lane 0's RNG and broadcast */
|
||||
float eps_out_j;
|
||||
if (lane_id == 0)
|
||||
eps_out_j = factorized_noise_fn(gpu_random_gaussian(rng));
|
||||
eps_out_j = __shfl_sync(0xFFFFFFFF, eps_out_j, 0);
|
||||
|
||||
const float* row = shmem_W + j * in_dim;
|
||||
float partial = 0.0f;
|
||||
for (int k = 0; k < local_count; k++) {
|
||||
int i = lane_id + k * 32;
|
||||
if (i < in_dim) {
|
||||
float eps_in_i = factorized_noise_fn(gpu_random_gaussian(rng));
|
||||
float w_noisy = row[i] + sigma_scale * eps_out_j * eps_in_i;
|
||||
partial += w_noisy * input_dist[k];
|
||||
}
|
||||
}
|
||||
/* Add noisy bias (lane 0 only) */
|
||||
if (lane_id == 0)
|
||||
partial += shmem_b[j] + sigma_scale * eps_out_j;
|
||||
/* Full warp reduction */
|
||||
float sum = warp_reduce_sum_all(partial);
|
||||
if (activate) sum = leaky_relu(sum);
|
||||
int gj = tile_offset + j;
|
||||
if (lane_id == (gj & 31))
|
||||
output_dist[gj >> 5] = sum;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Warp-cooperative matvec with broadcast output (non-distributed).
|
||||
*
|
||||
* Same as warp_matvec_leaky_relu_shmem but ALL lanes store every output
|
||||
* element (no strided scatter). Used for small output layers where all
|
||||
* lanes need the full result (e.g., C51 atom logits).
|
||||
*
|
||||
* @param output Non-distributed output array — all lanes write all elements
|
||||
*/
|
||||
__device__ void warp_matvec_broadcast_shmem(
|
||||
const float* __restrict__ shmem_W,
|
||||
const float* __restrict__ shmem_b,
|
||||
const float* input_dist,
|
||||
float* output, /* NOT distributed — all lanes get full copy */
|
||||
int in_dim,
|
||||
int out_dim,
|
||||
int tile_offset,
|
||||
int tile_rows,
|
||||
int activate,
|
||||
int lane_id
|
||||
) {
|
||||
(void)out_dim;
|
||||
for (int j = 0; j < tile_rows; j++) {
|
||||
const float* row = shmem_W + j * in_dim;
|
||||
float partial = 0.0f;
|
||||
for (int i = lane_id; i < in_dim; i += 32)
|
||||
partial += row[i] * input_dist[i / 32];
|
||||
if (lane_id == 0)
|
||||
partial += shmem_b[j];
|
||||
float sum = warp_reduce_sum_all(partial);
|
||||
if (activate) sum = leaky_relu(sum);
|
||||
/* Broadcast: ALL lanes store (redundant but correct for small arrays) */
|
||||
output[tile_offset + j] = sum;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Warp-cooperative noisy matvec with broadcast output.
|
||||
* Same as warp_noisy_matvec_leaky_relu_shmem but all lanes get full output.
|
||||
*/
|
||||
__device__ void warp_noisy_matvec_broadcast_shmem(
|
||||
const float* __restrict__ shmem_W,
|
||||
const float* __restrict__ shmem_b,
|
||||
const float* input_dist,
|
||||
float* output, /* NOT distributed — all lanes get full copy */
|
||||
int in_dim,
|
||||
int out_dim,
|
||||
int tile_offset,
|
||||
int tile_rows,
|
||||
int activate,
|
||||
float sigma_init,
|
||||
unsigned int* rng,
|
||||
int lane_id
|
||||
) {
|
||||
(void)out_dim;
|
||||
float sigma_scale = sigma_init / sqrtf((float)in_dim);
|
||||
int local_count = (in_dim + 31) / 32;
|
||||
|
||||
for (int j = 0; j < tile_rows; j++) {
|
||||
float eps_out_j;
|
||||
if (lane_id == 0)
|
||||
eps_out_j = factorized_noise_fn(gpu_random_gaussian(rng));
|
||||
eps_out_j = __shfl_sync(0xFFFFFFFF, eps_out_j, 0);
|
||||
|
||||
const float* row = shmem_W + j * in_dim;
|
||||
float partial = 0.0f;
|
||||
for (int k = 0; k < local_count; k++) {
|
||||
int i = lane_id + k * 32;
|
||||
if (i < in_dim) {
|
||||
float eps_in_i = factorized_noise_fn(gpu_random_gaussian(rng));
|
||||
float w_noisy = row[i] + sigma_scale * eps_out_j * eps_in_i;
|
||||
partial += w_noisy * input_dist[k];
|
||||
}
|
||||
}
|
||||
if (lane_id == 0)
|
||||
partial += shmem_b[j] + sigma_scale * eps_out_j;
|
||||
float sum = warp_reduce_sum_all(partial);
|
||||
if (activate) sum = leaky_relu(sum);
|
||||
output[tile_offset + j] = sum;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Warp-cooperative in-place RMSNorm on distributed data.
|
||||
*
|
||||
* Each lane holds elements in strided layout. Sum-of-squares is computed
|
||||
* via warp reduction, then each lane normalizes its own elements.
|
||||
*
|
||||
* @param data_dist Distributed data [DIST_SIZE(dim)] per lane (modified in place)
|
||||
* @param gamma Non-distributed gamma weights [dim] in global memory
|
||||
* @param dim Full dimension
|
||||
* @param lane_id Lane index (0..31)
|
||||
*/
|
||||
__device__ void warp_rmsnorm_inplace(float* data_dist, const float* __restrict__ gamma, int dim, int lane_id) {
|
||||
/* Each lane computes partial sum of squares over its elements */
|
||||
float partial_sq = 0.0f;
|
||||
for (int i = lane_id; i < dim; i += 32) {
|
||||
float val = data_dist[i / 32];
|
||||
partial_sq += val * val;
|
||||
}
|
||||
/* Warp reduction for total sum of squares */
|
||||
float total_sq = warp_reduce_sum_all(partial_sq);
|
||||
float rms = sqrtf(total_sq / (float)dim + 1e-6f);
|
||||
float inv_rms = 1.0f / rms;
|
||||
|
||||
/* Each lane normalizes its own elements */
|
||||
for (int i = lane_id; i < dim; i += 32) {
|
||||
data_dist[i / 32] = data_dist[i / 32] * inv_rms * gamma[i];
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* In-place RMSNorm: data[i] = (data[i] / RMS(data)) * gamma[i].
|
||||
*
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -197,6 +197,10 @@ pub struct GpuExperienceCollector {
|
||||
state_dim: usize,
|
||||
/// Network hidden dims — needed for shared memory tile sizing at launch time.
|
||||
network_dims: (usize, usize, usize, usize),
|
||||
/// Warp-cooperative kernel for sm_90+ (H100). None if not available.
|
||||
warp_kernel_func: Option<CudaFunction>,
|
||||
/// Whether to use warp-cooperative kernel (sm_90+ detected).
|
||||
use_warp_kernel: bool,
|
||||
/// Number of episodes buffers were allocated for (from config, not MAX constant).
|
||||
alloc_episodes: usize,
|
||||
/// Number of timesteps buffers were allocated for (from config, not MAX constant).
|
||||
@@ -295,7 +299,8 @@ impl GpuExperienceCollector {
|
||||
let full_source = format!("{common_src}\n{dim_overrides}\n{kernel_src}");
|
||||
info!(
|
||||
"GPU experience collector: compiling kernel with dims state={state_dim} market={market_dim} \
|
||||
shared=[{shared_h1},{shared_h2}] value={value_h} adv={adv_h} atoms_max={num_atoms_max}"
|
||||
shared=[{shared_h1},{shared_h2}] value={value_h} adv={adv_h} atoms_max={num_atoms_max} \
|
||||
shmem_max_in={shmem_max_in_dim}"
|
||||
);
|
||||
let ptx: Ptx = cudarc::nvrtc::compile_ptx(&full_source).map_err(|e| {
|
||||
MLError::ModelError(format!(
|
||||
@@ -315,6 +320,50 @@ impl GpuExperienceCollector {
|
||||
})?;
|
||||
info!("GPU experience collector: CUDA kernel compiled and loaded");
|
||||
|
||||
// Detect compute capability for kernel variant selection.
|
||||
// sm_90+ (H100/H200) uses warp-cooperative kernel for lower register pressure.
|
||||
let sm_major = {
|
||||
use cudarc::driver::sys::CUdevice_attribute;
|
||||
context
|
||||
.attribute(CUdevice_attribute::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR)
|
||||
.map_err(|e| {
|
||||
MLError::ModelError(format!("Failed to query compute capability: {e}"))
|
||||
})?
|
||||
};
|
||||
let use_warp_kernel = sm_major >= 9;
|
||||
|
||||
let warp_kernel_func = if use_warp_kernel {
|
||||
match module.load_function("dqn_full_experience_kernel_warp") {
|
||||
Ok(f) => {
|
||||
info!("GPU experience collector: warp-cooperative kernel loaded (sm_{sm_major}0+)");
|
||||
Some(f)
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"Warp kernel load failed on sm_{sm_major}0, falling back to standard: {e}"
|
||||
);
|
||||
None
|
||||
}
|
||||
}
|
||||
} else {
|
||||
info!("GPU experience collector: using standard kernel (sm_{sm_major}0 < sm_90)");
|
||||
None
|
||||
};
|
||||
|
||||
// Increase per-thread stack size for large hidden dimensions.
|
||||
// With SHARED_H1=512 (H100), scratch arrays consume ~5.5KB per thread,
|
||||
// exceeding the default 1KB CUDA stack limit.
|
||||
{
|
||||
use cudarc::driver::sys::{cuCtxSetLimit, CUlimit, CUresult};
|
||||
let stack_bytes = 16384_usize; // 16KB — enough for 5.5KB scratch + noise + call frames
|
||||
let result = unsafe { cuCtxSetLimit(CUlimit::CU_LIMIT_STACK_SIZE, stack_bytes) };
|
||||
if result != CUresult::CUDA_SUCCESS {
|
||||
tracing::warn!("cuCtxSetLimit(STACK_SIZE, {stack_bytes}) failed: {result:?}");
|
||||
} else {
|
||||
tracing::info!("CUDA stack size set to {stack_bytes} bytes");
|
||||
}
|
||||
}
|
||||
|
||||
// ---- Step 2: Extract weights ----
|
||||
let online_weights = extract_dueling_weights(online_vars, &stream)?;
|
||||
let target_weights = extract_dueling_weights(target_vars, &stream)?;
|
||||
@@ -469,6 +518,8 @@ impl GpuExperienceCollector {
|
||||
kernel_func,
|
||||
state_dim,
|
||||
network_dims,
|
||||
warp_kernel_func,
|
||||
use_warp_kernel,
|
||||
alloc_episodes,
|
||||
alloc_timesteps,
|
||||
online_weights,
|
||||
@@ -561,26 +612,42 @@ impl GpuExperienceCollector {
|
||||
* std::mem::size_of::<f32>() as u32;
|
||||
|
||||
let n = n_episodes as u32;
|
||||
// Force block size to 256 for shared memory cooperative loading
|
||||
let block_x = 256_u32;
|
||||
let grid_x = (n + block_x - 1) / block_x;
|
||||
let grid_dim = (grid_x, 1, 1);
|
||||
let block_dim = (block_x, 1, 1);
|
||||
let launch_config = LaunchConfig {
|
||||
grid_dim,
|
||||
block_dim,
|
||||
shared_mem_bytes: shmem_bytes,
|
||||
};
|
||||
|
||||
// Select kernel variant: warp-cooperative for sm_90+ (all modes including C51).
|
||||
let use_warp = self.use_warp_kernel
|
||||
&& self.warp_kernel_func.is_some();
|
||||
|
||||
let (launch_config, active_kernel) =
|
||||
if let (true, Some(warp_fn)) = (use_warp, &self.warp_kernel_func) {
|
||||
// Warp kernel: 1 warp (32 threads) per block, 1 block per episode.
|
||||
let cfg = LaunchConfig {
|
||||
grid_dim: (n, 1, 1),
|
||||
block_dim: (32, 1, 1),
|
||||
shared_mem_bytes: shmem_bytes,
|
||||
};
|
||||
(cfg, warp_fn)
|
||||
} else {
|
||||
// Standard kernel: 256 threads per block, 1 thread per episode.
|
||||
let block_x = 256_u32;
|
||||
let grid_x = (n + block_x - 1) / block_x;
|
||||
let cfg = LaunchConfig {
|
||||
grid_dim: (grid_x, 1, 1),
|
||||
block_dim: (block_x, 1, 1),
|
||||
shared_mem_bytes: shmem_bytes,
|
||||
};
|
||||
(cfg, &self.kernel_func)
|
||||
};
|
||||
|
||||
debug!(
|
||||
n_episodes,
|
||||
timesteps,
|
||||
grid_x = grid_dim.0,
|
||||
block_x = block_dim.0,
|
||||
grid_x = launch_config.grid_dim.0,
|
||||
block_x = launch_config.block_dim.0,
|
||||
shared_mem_kb = shmem_bytes / 1024,
|
||||
epsilon = config.epsilon,
|
||||
gamma = config.gamma,
|
||||
"Launching dqn_full_experience_kernel (shmem weight tiling)"
|
||||
warp_kernel = use_warp,
|
||||
"Launching dqn experience kernel"
|
||||
);
|
||||
|
||||
// ---- Step 5: Launch kernel ----
|
||||
@@ -590,9 +657,10 @@ impl GpuExperienceCollector {
|
||||
// Safety: kernel parameter order matches dqn_experience_kernel.cu lines 394-468 exactly.
|
||||
// All buffer sizes are pre-validated (MAX_EPISODES, MAX_TIMESTEPS).
|
||||
// All CudaSlice lifetimes are valid (owned by self).
|
||||
// Both standard and warp kernels share the same parameter signature.
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.kernel_func)
|
||||
.launch_builder(active_kernel)
|
||||
// Market data
|
||||
.arg(market_features_buf)
|
||||
.arg(targets_buf)
|
||||
@@ -680,7 +748,7 @@ impl GpuExperienceCollector {
|
||||
.arg(&mut self.td_error_out)
|
||||
.launch(launch_config)
|
||||
.map_err(|e| {
|
||||
MLError::ModelError(format!("dqn_full_experience_kernel launch failed: {e}"))
|
||||
MLError::ModelError(format!("dqn experience kernel launch failed: {e}"))
|
||||
})?;
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user