gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: stateless Philox for stochastic depth RNG (replace LCG)

Browse Source 
stochastic_depth_rng kernel now uses Philox hash(step, layer)
instead of stateful LCG. Monotonic step counter replaces
accumulated RNG state.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-13 18:39:46 +02:00
parent 6ad96712f3
commit 3177cde6d1
1 changed files with 12 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu
View File

@@ -1187,18 +1187,25 @@ extern "C" __global__ void bn_tanh_concat_f32_kernel(
* ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void stochastic_depth_rng(
float* __restrict__ scale_buf, /* [3] output: per-layer scales */
unsigned int* __restrict__ rng_state, /* [1] LCG state (read-write) */
unsigned int* __restrict__ rng_state, /* [1] step counter (read-write) */
float drop_prob /* probability of dropping each layer */
) {
unsigned int rng = rng_state[0];
/* Stateless Philox: deterministic per (step, layer). No accumulated LCG state. */
unsigned int step = rng_state[0];
float keep_scale = 1.0f / (1.0f - drop_prob);
for (int layer = 0; layer < 3; layer++) {
rng = rng * 1664525u + 1013904223u;
float u = (float)(rng & 0x00FFFFFFu) / 16777216.0f;
unsigned int key = step;
unsigned int ctr = (unsigned int)(layer * 37 + 7777);
ctr ^= key * 0x9E3779B9u;
ctr *= 0x85ebca6bu;
ctr ^= ctr >> 13;
ctr *= 0xc2b2ae35u;
ctr ^= ctr >> 16;
float u = (float)(ctr & 0x00FFFFFFu) / 16777216.0f;
scale_buf[layer] = (u > drop_prob) ? keep_scale : 0.0f;
}
rng_state[0] = rng;
rng_state[0] = step + 1; /* monotonic counter, not LCG */
}
/* ══════════════════════════════════════════════════════════════════════