dropout: device RNG kernel + Tensor fwd/bwd (T18)
csrc/ops/dropout.cu: counter-based RNG (splitmix64 over seed^index) → fp32 uniform → Bernoulli(keep=1-p); fwd writes out=x⊙mask + an fp32 mask buffer (per-element 1/(1-p) or 0); bwd applies the same mask (dx=d⊙mask). fp32 + bf16 activation variants (mask fp32 in both; uniform is dtype-independent so masks match across precisions). Stateless → re-run with same seed = same mask (T13 recompute-safe). Registered in build.rs + FFI decls. Tensor::dropout(p,seed)->(out,mask) and Tensor::dropout_backward(d,mask) wrap the launches (contiguous F32/BF16, default stream, per-op sync via the kernels). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
@@ -447,3 +447,48 @@ unsafe extern "C" {
|
||||
s: CudaStream,
|
||||
);
|
||||
}
|
||||
|
||||
// Dropout (Phase T18, csrc/ops/dropout.cu). A counter-based (stateless) RNG: the
|
||||
// keep/drop decision for element `i` is `hash(seed, i)` — no global state, so a
|
||||
// re-run with the same `seed` reproduces the same mask (compatible with T13
|
||||
// activation recomputation). Forward writes `out = x ⊙ mask` and the fp32 `mask`
|
||||
// buffer (mask[i] = (1/(1-p)) if kept else 0, the inverted-dropout scale);
|
||||
// backward applies the SAME mask: dx = d ⊙ mask. fp32 + bf16 activation variants
|
||||
// (mask is fp32 in both; the uniform is computed in fp32, dtype-independent).
|
||||
#[cfg(not(no_cuda))]
|
||||
unsafe extern "C" {
|
||||
pub fn launch_dropout_fwd_f32(
|
||||
x: *const f32,
|
||||
out: *mut f32,
|
||||
mask: *mut f32,
|
||||
p: f32,
|
||||
scale: f32,
|
||||
seed: u64,
|
||||
n: i32,
|
||||
s: CudaStream,
|
||||
);
|
||||
pub fn launch_dropout_bwd_f32(
|
||||
d: *const f32,
|
||||
mask: *const f32,
|
||||
dx: *mut f32,
|
||||
n: i32,
|
||||
s: CudaStream,
|
||||
);
|
||||
pub fn launch_dropout_fwd_bf16(
|
||||
x: *const c_void,
|
||||
out: *mut c_void,
|
||||
mask: *mut f32,
|
||||
p: f32,
|
||||
scale: f32,
|
||||
seed: u64,
|
||||
n: i32,
|
||||
s: CudaStream,
|
||||
);
|
||||
pub fn launch_dropout_bwd_bf16(
|
||||
d: *const c_void,
|
||||
mask: *const f32,
|
||||
dx: *mut c_void,
|
||||
n: i32,
|
||||
s: CudaStream,
|
||||
);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user