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

fix: VSN softmax → sigmoid — prevents 256× signal collapse

Browse Source 
Softmax over SH2=256 features gives mask[i]=1/256=0.004 per feature,
collapsing the input signal 256×. Sigmoid gives mask[i]∈(0,1) independently
per feature. At W_vsn2=0 init: sigmoid(0)=0.5 → h_masked = 0.5*h_s2.
Shared memory reduced: SH2+R (no mask buffer needed, sigmoid is inline).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-15 08:04:08 +02:00
parent 64898b0cce
commit 291e2d07d7
2 changed files with 12 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
crates/ml/src/cuda_pipeline/batched_forward.rs
View File

@@ -964,7 +964,7 @@ impl CublasGemmSet {
// ── 1. VSN bottleneck: h_s2 [B, SH2] → vsn_masked [B, SH2] ──
let w_vsn1_idx = 26 + d * 2; // W_vsn1_d
let w_vsn2_idx = 26 + d * 2 + 1; // W_vsn2_d
let smem = (sh2 + r + sh2) * std::mem::size_of::<f32>();
let smem = (sh2 + r) * std::mem::size_of::<f32>(); // h_local[SH2] + proj[R], no mask buf (sigmoid is inline)
unsafe {
stream
.launch_builder(vsn_k)

54
crates/ml/src/cuda_pipeline/experience_kernels.cu
View File

@@ -2643,16 +2643,20 @@ extern "C" __global__ void strided_scatter(
/* ================================================================== */
/**
* TFT-inspired per-branch feature selection with R=16 bottleneck.
* Per-branch feature gating with R=16 bottleneck.
*
* h_masked[b,i] = h_s2[b,i] * softmax(h_s2[b] @ W_vsn1 @ W_vsn2)[i]
* h_masked[b,i] = h_s2[b,i] * sigmoid(h_s2[b] @ W_vsn1 @ W_vsn2)[i]
*
* Sigmoid (not softmax): each feature is independently gated ∈ (0,1).
* At W_vsn2=0 init: sigmoid(0) = 0.5 → h_masked = 0.5 * h_s2 (stable start).
* Softmax would give 1/SH2 ≈ 0.004 per feature → 256× signal collapse → gradient explosion.
*
* Phase 1: proj[r] = sum_k h_local[k] * W_vsn1[k*R + r]
* Phase 2: mask[i] = sum_r proj[r] * W_vsn2[r*SH2 + i], then softmax
* Phase 2: mask[i] = sigmoid(sum_r proj[r] * W_vsn2[r*SH2 + i])
* Phase 3: h_masked[b,i] = h_local[i] * mask[i]
*
* Grid: (B, 1, 1), Block: (256, 1, 1)
* Dynamic shared: (SH2 + R + SH2) * sizeof(float)
* Dynamic shared: (SH2 + R) * sizeof(float)
*/
extern "C" __global__ void variable_select_bottleneck(
const float* __restrict__ h_s2, /* [B, SH2] */
@@ -2664,7 +2668,6 @@ extern "C" __global__ void variable_select_bottleneck(
extern __shared__ float smem[];
float* h_local = smem; /* [SH2] */
float* proj = smem + SH2; /* [R] */
float* mask = smem + SH2 + R; /* [SH2] */
int b = blockIdx.x;
if (b >= B) return;
@@ -2685,49 +2688,14 @@ extern "C" __global__ void variable_select_bottleneck(
}
__syncthreads();
/* Phase 2: mask[i] = sum_r proj[r] * W_vsn2[r*SH2 + i], then softmax */
/* Phase 2+3: mask[i] = sigmoid(proj @ W_vsn2[:,i]), h_masked = h_local * mask */
for (int i = threadIdx.x; i < SH2; i += blockDim.x) {
float acc = 0.0f;
for (int r = 0; r < R; r++) {
acc += proj[r] * W_vsn2[r * SH2 + i];
}
mask[i] = acc;
}
__syncthreads();
/* Softmax over mask[0..SH2) — find max, then sum_exp, then normalize */
/* Parallel reduction for max */
__shared__ float red_buf[256];
float my_max = -1e38f;
for (int i = threadIdx.x; i < SH2; i += blockDim.x) {
my_max = fmaxf(my_max, mask[i]);
}
red_buf[threadIdx.x] = my_max;
__syncthreads();
for (int s = blockDim.x / 2; s > 0; s >>= 1) {
if (threadIdx.x < s) red_buf[threadIdx.x] = fmaxf(red_buf[threadIdx.x], red_buf[threadIdx.x + s]);
__syncthreads();
}
float g_max = red_buf[0];
/* Compute exp and store back, then sum */
float my_sum = 0.0f;
for (int i = threadIdx.x; i < SH2; i += blockDim.x) {
float e = expf(mask[i] - g_max);
mask[i] = e;
my_sum += e;
}
red_buf[threadIdx.x] = my_sum;
__syncthreads();
for (int s = blockDim.x / 2; s > 0; s >>= 1) {
if (threadIdx.x < s) red_buf[threadIdx.x] += red_buf[threadIdx.x + s];
__syncthreads();
}
float g_sum = fmaxf(red_buf[0], 1e-8f);
/* Phase 3: h_masked[b,i] = h_local[i] * (mask[i] / g_sum) */
for (int i = threadIdx.x; i < SH2; i += blockDim.x) {
h_masked[b * SH2 + i] = h_local[i] * (mask[i] / g_sum);
float gate = 1.0f / (1.0f + expf(-acc)); /* sigmoid per feature */
h_masked[b * SH2 + i] = h_local[i] * gate;
}
}