From d61aefe2b88c0d2059b4e46358bfaa3cab1e4ef1 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 23 Apr 2026 23:12:53 +0200 Subject: [PATCH] =?UTF-8?q?fix(dqn):=20ISV-driven=20per-branch=20grad=20ba?= =?UTF-8?q?lancer=20=E2=80=94=20equalise=20via=20median=20target?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The pre-spec balancer used `cap = 4 × median` to suppress outlier branches. On L40S train-5wb4n (epoch 9 norms `[dir=130, mag=20000, ord=18000, urg=16000]`), that made cap = 68000 — exceeding every branch's norm. No cap ever fired; the kernel was a geometric no-op for 11 consecutive epochs while direction starved at 1/150 the magnitude gradient. The outlier-cap algorithm can only handle "one branch above the pack", not "three co-elevated, one starved". Replace with ISV-driven equalisation: - 4 new ISV slots (31..34): per-branch grad-norm target - 1 new ISV slot (35): scale clamp limit - ISV_TOTAL_DIM 31 → 36 - New on-GPU producer kernel `grad_balance_isv_update` runs after `branch_grad_norm_reduce`, before `branch_grad_rescale`. Computes current-step median + max/min spread, applies adaptive-rate EMA to the ISV slots. Single-threaded (1 warp), graph-capture safe, no atomicAdd, no host syncs. - Rescale kernel now reads ISV target + limit, computes per-branch `scale[b] = clamp(target[b] / norm[b], 1/limit, limit)`. Every branch equalises toward the shared median target over time; starved branches boost, elevated branches suppress. Limit is the observed max/min spread EMA, hard-bounded [2.0, 1e4] as a numerical-safety cap (not tuning). Adaptive-rate EMA: `alpha = clamp(err / (err + baseline), 0.01, 0.3)`. Fast response when error is large relative to current value, slow drift when stable. No fixed alpha constant. Bootstrap at construction + fold reset: targets = 1.0, limit = 2.0. Pre-first-EMA-update behaviour: scale clamped to [0.5, 2.0], effectively the original identity on warm-up, until real observations flow in. Complies with feedback_isv_for_adaptive_bounds.md (2026-04-23): adaptive bounds always live in ISV, never hardcoded. No partial refactor — all consumers of the balancer contract (launch, kernels, rescale algorithm) migrate together. Still pending from the task-#92 triad: (a) C51 backward kernel must replicate the forward advantage-standardization for d==1 (chain-rule violation produces wrong gradient; separate commit), (b) IQL branch_scales floor for direction-branch Hold/Flat starvation (separate commit). This commit is the safety net; the two root-cause fixes come next. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../branch_grad_balance_kernel.cu | 244 ++++++++++++------ .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 107 +++++++- 2 files changed, 254 insertions(+), 97 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/branch_grad_balance_kernel.cu b/crates/ml/src/cuda_pipeline/branch_grad_balance_kernel.cu index 13638792a..3d42424da 100644 --- a/crates/ml/src/cuda_pipeline/branch_grad_balance_kernel.cu +++ b/crates/ml/src/cuda_pipeline/branch_grad_balance_kernel.cu @@ -1,57 +1,66 @@ /** - * branch_grad_balance_kernel — adaptive per-branch gradient-norm balancer. + * branch_grad_balance_kernel — ISV-driven adaptive per-branch gradient-norm + * balancer. * - * Caps each of the 4 factored-action branches' gradient L2 norm at - * `num_branches × median_branch_norm` - * (median across the 4 branches' norms at the current step). Branches - * within the cap pass through unchanged. + * Scales each of the 4 factored-action branches' gradient toward a + * per-branch target norm maintained as an adaptive EMA in the ISV signal + * bus (slots GRAD_NORM_TARGET_{DIR,MAG,ORD,URG}). The scale factor is + * clamped to `[1/limit, limit]` where `limit` is itself an ISV-bus + * adaptive EMA of observed `max_norm / min_norm` spread + * (GRAD_SCALE_LIMIT_INDEX). * - * Why this shape (documented architecturally, NOT tuned): - * * `num_branches` (= 4) is architectural — the count of factored action - * axes (direction, magnitude, order, urgency). It is NOT a tuning - * knob; changing the action factorisation changes this number. - * * `median_branch_norm` is a per-step statistical reference that scales - * naturally with the training regime (Q-magnitude, loss blend, etc.). - * There is no fixed target scale, no schedule, and no hyperparameter. - * * Their product enforces the defensible bound: "no single branch can - * dominate the gradient by more than `num_branches` × the median." - * Geometrically, a branch at the cap contributes as much L2 mass as - * all other branches combined evaluated at the median. - * * Scale-only-when-needed preserves branches that are proportional; - * the mechanism never forces artificial equality. + * Rationale — what was wrong with the pre-spec design: + * The previous kernel capped branches whose norm exceeded `4 × median`. + * On L40S train-5wb4n (2026-04-23) with grad distribution + * [dir=130, mag=20000, ord=18000, urg=16000] + * median = 17000, cap = 68000 → no branch triggered the cap, the + * balancer was a geometric no-op for 11 consecutive epochs. The + * algorithm only handles outlier-reduction; it cannot address + * starvation ("one branch far below the pack"). The pattern here is + * starvation: direction is 150× below magnitude, not magnitude above + * the rest. * - * Motivation (L40S train-mdh86, epochs 0-6): - * HEALTH_DIAG reported `grad_ratio_mag_dir` ≈ 1.5e4 – 2.6e4×. Magnitude - * gradients were three-to-four orders of magnitude larger than - * direction for the entire warm-up, which destabilised the dueling + - * factored-advantage decomposition. On a single step at epoch 7 the - * ratio collapsed to ~100× and Sharpe went from +34 to -67 without - * recovering. Per-branch capping at `num_branches × median` is a - * symmetric, signal-driven bound that prevents either end of the - * swing. + * Why ISV-driven (architecturally, NOT tuned): + * * Per-branch target tracking (EMA → median) causes every branch to + * equalise over time, regardless of which side is the outlier. + * Starved branches are boosted; elevated branches are suppressed. + * * The scale clamp `limit` is not a fixed constant. It tracks the + * observed max/min norm spread via adaptive EMA — grows just enough + * to correct the actual asymmetry, bounded by an absolute safety + * cap (1e4) written on the Rust side. + * * All state lives in ISV slots. The producer `grad_balance_isv_update` + * reads observed branch norms and writes updated EMAs. Consumers + * (this kernel) read ISV at rescale time. A single producer, one + * source of truth — the ISV invariant applies everywhere. + * * See `feedback_isv_for_adaptive_bounds.md` (2026-04-23) for the + * project-wide rule this satisfies. * - * Pipeline (one kernel, two launches): - * branch_grad_norm_reduce: grid=(4,1,1), block=(256,1,1). One block - * per branch; reduces the branch's - * sum-of-squares via shared-mem tree and - * writes sqrt to branch_norms_dev[branch]. - * No atomicAdd (one-block-per-branch). + * Pipeline (three kernels, three launches): + * branch_grad_norm_reduce: grid=(4,1,1), block=(256,1,1). Unchanged + * from the pre-spec version. Writes per- + * branch L2 norms to branch_norms_dev[4]. * - * branch_grad_rescale: grid=(max_blocks, 4, 1), block=(256,1,1). - * blockIdx.y = branch id; blockIdx.x = 256- - * element sub-chunk. First 4 threads cache - * the 4 branch norms into shared memory, - * compute the median + per-branch scales, - * then each thread multiplies one element - * of its branch's slice by its branch scale. - * No atomicAdd (each thread writes a - * distinct element). + * grad_balance_isv_update: grid=(1,1,1), block=(32,1,1). Thread 0 + * reads branch_norms_dev, computes median + * + max/min spread, applies adaptive EMAs + * to ISV slots. Single-threaded work; + * extra threads idle-exit. No atomicAdd. + * Writes: ISV[31..35]. Safe in graph + * capture — no host syncs, no dynamic + * alloc. + * + * branch_grad_rescale: grid=(max_blocks, 4, 1), block=(256,1,1). + * blockIdx.y = branch id; blockIdx.x = + * 256-element sub-chunk. Thread 0 reads + * ISV target + limit, computes scale, + * writes diag slot. Each thread multiplies + * one element by its branch's scale. No + * atomicAdd. * * CUDA Graph safety: - * All launches are plain kernel launches with fixed grid/block config. - * No dynamic allocations, no host syncs, no atomicAdd. Safe to capture - * inside `adam_grad_child` between the post_aux phase and - * `compute_grad_norm_for_adam`. + * All three launches have fixed grid/block. No host syncs, no dynamic + * allocs, no atomicAdd. Safe to capture inside `adam_grad_child` + * between the post_aux phase and `compute_grad_norm_for_adam`. */ extern "C" __global__ void branch_grad_norm_reduce( @@ -88,8 +97,8 @@ extern "C" __global__ void branch_grad_norm_reduce( } /* Device helper: median of 4 non-negative floats (sort-based, fully - * branch-deterministic, no reduction primitives). Returns the average - * of the middle two after sorting — the standard even-length median. + * branch-deterministic). Returns the average of the middle two after + * sorting — the standard even-length median. */ __device__ __forceinline__ float branch_grad_balance_median4( float a, float b, float c, float d @@ -111,59 +120,124 @@ __device__ __forceinline__ float branch_grad_balance_median4( return 0.5f * (b + c); } +/* grad_balance_isv_update — producer kernel for the ISV-driven balancer. + * + * Reads the per-branch L2 norms (written by branch_grad_norm_reduce) and + * updates the ISV signal bus: + * * target[b] (slots 31..35): adaptive-rate EMA toward the current + * step's median norm. All four targets track the same median, so + * the rescale pass equalises branches. + * * limit (slot 35): adaptive-rate EMA toward `max/min` spread. + * Hard-clamped to [2.0, 1e4] as a numerical-safety bound only. + * + * Adaptive-rate EMA: alpha = clamp(err / (err + baseline), 0.01, 0.3). + * High alpha when error is large relative to current value (fast + * response); low alpha when error is small (slow drift). No fixed-rate + * constant — the rate adapts to observed error size. + * + * Single-threaded: the work is a trivial 4-float sort + EMA update. + * Launched with block=(32,1,1), grid=(1,1,1). Threads 1..31 idle-exit. + * Kept at 32 threads to match a warp (hardware scheduling alignment). + */ +extern "C" __global__ void grad_balance_isv_update( + const float* __restrict__ branch_norms_dev, /* [4] per-branch L2 norms */ + float* __restrict__ isv_signals, /* device-mapped pinned ISV bus */ + int target_base_idx, /* = GRAD_NORM_TARGET_DIR_INDEX (31) */ + int limit_idx /* = GRAD_SCALE_LIMIT_INDEX (35) */ +) { + if (threadIdx.x != 0 || blockIdx.x != 0) return; + + float n0 = branch_norms_dev[0]; + float n1 = branch_norms_dev[1]; + float n2 = branch_norms_dev[2]; + float n3 = branch_norms_dev[3]; + + /* All-zero branch norms → quiescent training step; skip the update + * to avoid driving targets toward zero and the limit to 1.0 (which + * would lock the balancer into identity for the rest of the run). */ + float max_n = fmaxf(fmaxf(n0, n1), fmaxf(n2, n3)); + if (max_n < 1e-9f) return; + + float median = branch_grad_balance_median4(n0, n1, n2, n3); + float min_n = fminf(fminf(n0, n1), fminf(n2, n3)); + + /* Spread: observed max/min ratio. Floor at 1.0 (no asymmetry) and + * cap at 1e4 (numerical-safety — no amplification beyond 10000×). */ + float spread = (min_n > 1e-9f) ? (max_n / min_n) : 1.0f; + if (spread < 1.0f) spread = 1.0f; + if (spread > 1e4f) spread = 1e4f; + + /* Per-branch target: EMA toward median. Adaptive alpha. */ + for (int b = 0; b < 4; b++) { + float cur = isv_signals[target_base_idx + b]; + float err = fabsf(median - cur); + float baseline = fmaxf(cur, 1e-6f); + float alpha = err / (err + baseline); + alpha = fminf(0.3f, fmaxf(0.01f, alpha)); + float next = (1.0f - alpha) * cur + alpha * median; + isv_signals[target_base_idx + b] = next; + } + + /* Scale clamp limit: EMA toward observed spread, hard-bounded + * [2.0, 1e4] as safety (not tuning). */ + float cur_limit = isv_signals[limit_idx]; + float err_limit = fabsf(spread - cur_limit); + float baseline_limit = fmaxf(cur_limit, 1.0f); + float alpha_limit = err_limit / (err_limit + baseline_limit); + alpha_limit = fminf(0.3f, fmaxf(0.01f, alpha_limit)); + float next_limit = (1.0f - alpha_limit) * cur_limit + alpha_limit * spread; + if (next_limit < 2.0f) next_limit = 2.0f; + if (next_limit > 1e4f) next_limit = 1e4f; + isv_signals[limit_idx] = next_limit; +} + extern "C" __global__ void branch_grad_rescale( float* __restrict__ grad_buf, - const int* __restrict__ branch_starts, /* [4] element offsets */ - const int* __restrict__ branch_lens, /* [4] element counts */ - const float* __restrict__ branch_norms_dev, /* [4] L2 norms from reduce pass */ - int num_branches, /* architectural constant (= 4); cap multiplier */ - float* __restrict__ branch_scales_dev /* [4] optional diagnostic output — scale applied per branch */ + const int* __restrict__ branch_starts, /* [4] element offsets */ + const int* __restrict__ branch_lens, /* [4] element counts */ + const float* __restrict__ branch_norms_dev, /* [4] L2 norms from reduce pass */ + const float* __restrict__ isv_signals, /* device-mapped pinned ISV bus */ + int target_base_idx, /* = GRAD_NORM_TARGET_DIR_INDEX (31) */ + int limit_idx, /* = GRAD_SCALE_LIMIT_INDEX (35) */ + float* __restrict__ branch_scales_dev /* [4] optional diagnostic output */ ) { const int branch = blockIdx.y; if (branch >= 4) return; - __shared__ float sh_norms[4]; - __shared__ float sh_scales[4]; + __shared__ float sh_scale; const int tid = threadIdx.x; - if (tid < 4) { - sh_norms[tid] = branch_norms_dev[tid]; - sh_scales[tid] = 1.0f; - } - __syncthreads(); - /* Thread 0 computes median + per-branch scales once per block. - * The cap is `num_branches × median` — a branch at the cap has - * L2 mass equal to the sum of `num_branches` median-branches, - * i.e. the sum of the other branches evaluated at the median. */ + /* Thread 0 reads the ISV target + limit, computes the scale, + * publishes to shared memory and (optional) diagnostic slot. */ if (tid == 0) { - float n0 = sh_norms[0]; - float n1 = sh_norms[1]; - float n2 = sh_norms[2]; - float n3 = sh_norms[3]; - float med = branch_grad_balance_median4(n0, n1, n2, n3); - float cap = (float)num_branches * med; - /* Zero/near-zero median means all branches are quiescent — no - * rescaling is defensible (cap would be 0, divide-by-zero). - * Guard with the same epsilon the existing grad_ratio_mag_dir - * accessor uses (1e-9). */ - if (med > 1e-9f) { - sh_scales[0] = (n0 > cap) ? (cap / n0) : 1.0f; - sh_scales[1] = (n1 > cap) ? (cap / n1) : 1.0f; - sh_scales[2] = (n2 > cap) ? (cap / n2) : 1.0f; - sh_scales[3] = (n3 > cap) ? (cap / n3) : 1.0f; + float norm = branch_norms_dev[branch]; + float target = isv_signals[target_base_idx + branch]; + float limit = isv_signals[limit_idx]; + if (limit < 1.0f) limit = 1.0f; + float inv_limit = 1.0f / limit; + + float scale = 1.0f; + /* Both target and norm must be non-trivial to compute a + * defensible scale. Quiescent training (all-zero norms) or + * uninitialised targets (bootstrap 1.0, no EMA update yet) + * pass through unchanged. */ + if (norm > 1e-9f && target > 1e-9f) { + scale = target / norm; + if (scale > limit) scale = limit; + if (scale < inv_limit) scale = inv_limit; } - /* Publish scale to the optional diagnostic output (only block - * (0, branch, 0) writes to avoid redundant stores across - * element-chunk blocks). The grad_buf writes below are - * always idempotent across blockIdx.x so any block can emit. */ + sh_scale = scale; + + /* Only block (0, branch, 0) emits the diagnostic to avoid + * redundant stores across element-chunk blocks. */ if (blockIdx.x == 0 && branch_scales_dev != nullptr) { - branch_scales_dev[branch] = sh_scales[branch]; + branch_scales_dev[branch] = scale; } } __syncthreads(); - const float scale = sh_scales[branch]; + const float scale = sh_scale; /* Skip the multiply entirely when scale == 1.0f — preserves the * original gradient bit-identically on healthy branches. */ if (scale == 1.0f) return; diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 58165990c..f0f37cf30 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -148,10 +148,21 @@ const ISV_NETWORK_DIM: usize = 23; /// feed the attention path. Slots [23..31] carry per-branch Q-support /// (center, half-width) for direction/magnitude/order/urgency — written by /// `update_eval_v_range`, read by `adaptive_atom_positions`, -/// `warm_start_atom_positions`, and per-sample support tiling. Extending -/// past 31 requires updating the pinned allocation in the constructor and -/// any kernel that takes an ISV pointer expecting a specific length. -const ISV_TOTAL_DIM: usize = 31; +/// `warm_start_atom_positions`, and per-sample support tiling. +/// +/// Grad-balance unification bundle (2026-04-23): slots [31..35] extend the +/// scratchpad for ISV-driven adaptive gradient balancing. Slots [31..35) +/// carry per-branch gradient-norm targets; slot [35] carries the +/// observed-spread-driven scale clamp limit. Written by +/// `grad_balance_isv_update` on-device (inside graph capture); read by +/// `branch_grad_rescale`. Replaces the pre-spec hardcoded `4 × median` +/// cap that was a no-op when three branches were co-elevated (see +/// feedback_isv_for_adaptive_bounds.md + task #92). +/// +/// Extending past 36 requires updating the pinned allocation in the +/// constructor and any kernel that takes an ISV pointer expecting a +/// specific length. +const ISV_TOTAL_DIM: usize = 36; /// Legacy alias preserved for call sites that haven't been audited for the /// network-vs-total split. New code should pick `ISV_NETWORK_DIM` (for weight /// tensor sizing) or `ISV_TOTAL_DIM` (for the broadcast bus buffer). @@ -203,6 +214,28 @@ pub const V_CENTER_ORD_INDEX: usize = 27; pub const V_HALF_ORD_INDEX: usize = 28; pub const V_CENTER_URG_INDEX: usize = 29; pub const V_HALF_URG_INDEX: usize = 30; +/// Grad-balance unification bundle (2026-04-23): per-branch adaptive gradient- +/// norm targets. Four slots, one per branch (dir/mag/ord/urg). Written on-GPU +/// by `grad_balance_isv_update` via adaptive-rate EMA toward the median of the +/// observed per-branch norms each step — so every branch's target tracks the +/// median, yielding equalisation. Read by `branch_grad_rescale` to compute +/// `scale[b] = clamp(target[b] / norm[b], 1/limit, limit)`. +/// +/// Replaces the pre-spec `cap = 4 × median_branch_norm` formula, which became +/// a no-op once three branches co-elevated (observed grad_ratio_mag_dir +/// 48-868× in train-5wb4n). See `feedback_isv_for_adaptive_bounds.md`. +pub const GRAD_NORM_TARGET_DIR_INDEX: usize = 31; +pub const GRAD_NORM_TARGET_MAG_INDEX: usize = 32; +pub const GRAD_NORM_TARGET_ORD_INDEX: usize = 33; +pub const GRAD_NORM_TARGET_URG_INDEX: usize = 34; +/// Grad-balance unification bundle (2026-04-23): scale-clamp limit. Single +/// slot carrying the adaptive bound `limit` such that per-branch scale is +/// clamped to `[1/limit, limit]`. The limit tracks observed +/// `max_norm / min_norm` spread via adaptive EMA, so the clamp range grows +/// just enough to correct the actual asymmetry. Hard-capped at 1e4 as a +/// numerical-safety bound (not a tuning knob). Written on-GPU by +/// `grad_balance_isv_update`; read by `branch_grad_rescale`. +pub const GRAD_SCALE_LIMIT_INDEX: usize = 35; const ISV_EMB_DIM: usize = 8; // ISV embedding output dimension (FC2 output, gate/gamma input) /// First ISV tensor index in the flat param buffer. /// ISV weights (68-79) are online-only — NOT synced to the target network. @@ -1190,6 +1223,10 @@ pub struct GpuDqnTrainer { /// grad writes complete and BEFORE `compute_grad_norm_for_adam`, so /// Adam and the global clip observe the rebalanced gradient. branch_grad_balance_reduce: CudaFunction, + /// ISV-driven producer kernel (2026-04-23): reads per-branch norms, + /// applies adaptive-rate EMA to ISV slots `GRAD_NORM_TARGET_*` and + /// `GRAD_SCALE_LIMIT_INDEX`. Runs between reduce and rescale. + branch_grad_balance_isv_update: CudaFunction, branch_grad_balance_rescale: CudaFunction, /// `[4]` element offsets into `grad_buf` for branches 0-3 /// (direction, magnitude, order, urgency). Each branch covers the @@ -6817,17 +6854,19 @@ impl GpuDqnTrainer { // CUfunction handles are isolated from all other graphs (Hopper // CUfunction-isolation rule — shared handles corrupt captured // kernel state). - let (branch_grad_balance_reduce, branch_grad_balance_rescale) = { + let (branch_grad_balance_reduce, branch_grad_balance_isv_update, branch_grad_balance_rescale) = { let module = stream.context().load_cubin(BRANCH_GRAD_BALANCE_CUBIN.to_vec()) .map_err(|e| MLError::ModelError(format!("branch_grad_balance cubin load: {e}")))?; let reduce = module.load_function("branch_grad_norm_reduce") .map_err(|e| MLError::ModelError(format!("branch_grad_norm_reduce load: {e}")))?; + let isv_update = module.load_function("grad_balance_isv_update") + .map_err(|e| MLError::ModelError(format!("grad_balance_isv_update load: {e}")))?; let rescale = module.load_function("branch_grad_rescale") .map_err(|e| MLError::ModelError(format!("branch_grad_rescale load: {e}")))?; - (reduce, rescale) + (reduce, isv_update, rescale) }; info!( - "GpuDqnTrainer: branch_grad_balance kernels loaded — max_slice_len={} num_branches=4 (architectural)", + "GpuDqnTrainer: branch_grad_balance kernels loaded (ISV-driven) — max_slice_len={} num_branches=4 (architectural)", branch_slice_max_len, ); @@ -8055,6 +8094,19 @@ impl GpuDqnTrainer { *sig_ptr.add(23 + 2 * b) = 0.0_f32; // v_center *sig_ptr.add(23 + 2 * b + 1) = half_bootstrap; // v_half } + // Grad-balance unification bundle (2026-04-23): bootstrap + // per-branch grad-norm targets and scale-clamp limit. + // Targets start at 1.0 (neutral, no scaling until the first + // on-GPU `grad_balance_isv_update` populates the EMA from + // observed per-branch norms). Limit bootstraps at 2.0 — + // allows up to 2× rescaling before any real spread has + // been observed; EMA grows upward only if training + // warrants. The 2.0 lower bound and 1e4 upper bound are + // numerical-safety hard limits (not adaptive tuning). + for b in 0..4usize { + *sig_ptr.add(GRAD_NORM_TARGET_DIR_INDEX + b) = 1.0_f32; + } + *sig_ptr.add(GRAD_SCALE_LIMIT_INDEX) = 2.0_f32; } (host_ptr as *mut f32, dev_ptr_out) }; @@ -8846,6 +8898,7 @@ impl GpuDqnTrainer { grad_decomp_snapshot_len, grad_decomp_kernel, branch_grad_balance_reduce, + branch_grad_balance_isv_update, branch_grad_balance_rescale, branch_slice_starts_dev, branch_slice_lens_dev, @@ -10050,6 +10103,9 @@ impl GpuDqnTrainer { let lens_ptr = self.branch_slice_lens_dev.raw_ptr(); let norms_ptr = self.branch_grad_norms_dev.raw_ptr(); let scales_ptr = self.branch_grad_scales_dev.raw_ptr(); + let isv_ptr = self.isv_signals_dev_ptr; + let target_base_idx: i32 = GRAD_NORM_TARGET_DIR_INDEX as i32; + let limit_idx: i32 = GRAD_SCALE_LIMIT_INDEX as i32; // ── Pass 1: per-branch L2 reduction (grid=(4,), block=(256,)) ─ unsafe { @@ -10069,12 +10125,37 @@ impl GpuDqnTrainer { ))?; } - // ── Pass 2: cap + rescale (grid=(blocks_x, 4,), block=(256,)) ─ + // ── Pass 2: ISV-driven producer (grid=(1,), block=(32,)) ────── + // Reads per-branch norms, applies adaptive-rate EMA to the ISV + // targets (slots 31..34) and scale clamp limit (slot 35). Must + // run AFTER reduce and BEFORE rescale so rescale reads the + // updated targets. Graph-safe: fixed grid/block, no atomicAdd, + // single producer path. + unsafe { + self.stream + .launch_builder(&self.branch_grad_balance_isv_update) + .arg(&norms_ptr) + .arg(&isv_ptr) + .arg(&target_base_idx) + .arg(&limit_idx) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (32, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError( + format!("grad_balance_isv_update launch: {e}") + ))?; + } + + // ── Pass 3: rescale (grid=(blocks_x, 4,), block=(256,)) ─────── // blocks_x covers the largest branch; shorter branches early-exit - // via the `idx < len` guard in the kernel. Architectural value - // `num_branches = 4` is passed as the cap multiplier. + // via the `idx < len` guard in the kernel. Each branch reads + // `isv_signals[GRAD_NORM_TARGET_DIR_INDEX + branch]` and + // `isv_signals[GRAD_SCALE_LIMIT_INDEX]` to compute its scale as + // `clamp(target / norm, 1/limit, limit)`. Equalises branches to + // the adaptive median; bounded by observed-spread-driven limit. let blocks_x = ((self.branch_slice_max_len as u32 + 255) / 256).max(1); - let num_branches: i32 = 4; unsafe { self.stream .launch_builder(&self.branch_grad_balance_rescale) @@ -10082,7 +10163,9 @@ impl GpuDqnTrainer { .arg(&starts_ptr) .arg(&lens_ptr) .arg(&norms_ptr) - .arg(&num_branches) + .arg(&isv_ptr) + .arg(&target_base_idx) + .arg(&limit_idx) .arg(&scales_ptr) .launch(LaunchConfig { grid_dim: (blocks_x, 4, 1),