From 1f910d0aac83b21ca8599ea2d72aba4e02845eb9 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 9 Apr 2026 21:56:38 +0200 Subject: [PATCH] fix(critical): c51_grad writes branch-major d_adv_logits layout The old sample-major interleaved layout did not match cuBLAS backward branch-major pointer arithmetic. At batch=16384, backward GEMM read cross-contaminated data producing zero weight gradients. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml/src/cuda_pipeline/c51_grad_kernel.cu | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu b/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu index ec4fa71b9..a76b9da56 100644 --- a/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu @@ -78,9 +78,12 @@ extern "C" __global__ void c51_grad_kernel( int A_d = branch_sizes[d]; float inv_A = 1.0f / (float)A_d; - int branch_offset = 0; + /* Branch-major base offset: all B samples for branches 0..d-1 precede branch d. + * Layout: [B*B0*NA | B*B1*NA | B*B2*NA | B*B3*NA] (contiguous per-branch blocks). + * Matches cuBLAS backward pointer arithmetic for per-branch dY matrices. */ + int branch_base = 0; for (int dd = 0; dd < d; dd++) - branch_offset += branch_sizes[dd] * num_atoms; + branch_base += batch_size * branch_sizes[dd] * num_atoms; /* Per-branch loss weighting: magnitude (d==1) gets ZERO C51 gradient. * C51 cross-entropy inherently prefers low-variance actions (Small positions @@ -97,7 +100,7 @@ extern "C" __global__ void c51_grad_kernel( for (int a = 0; a < A_d; a++) { float dueling_grad = (a == a_d) ? (1.0f - inv_A) : (-inv_A); float grad_val = branch_scale * d_combined * dueling_grad; - int adv_idx = b * total_branch_atoms + branch_offset + a * num_atoms + j; + int adv_idx = branch_base + b * (A_d * num_atoms) + a * num_atoms + j; atomicAdd(&d_adv_logits[adv_idx], grad_val); } }