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

fix(dqn): ISV-driven per-branch grad balancer — equalise via median target

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-23 23:12:53 +02:00
parent 69c663d360
commit d61aefe2b8
2 changed files with 254 additions and 97 deletions
Download Patch File Download Diff File Expand all files Collapse all files

244
crates/ml/src/cuda_pipeline/branch_grad_balance_kernel.cu
View File

@@ -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;

107
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -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),