fix: IQN 4-branch runtime + order counterfactual + per-branch entropy + histogram
IQN kernel rewrite (critical — was corrupting 75% of gradient budget): - Eliminate compile-time BRANCH_*_SIZE defines (BRANCH_0_SIZE was 5, should be 3) - All 9 IQN kernels now take runtime b0/b1/b2/b3 params - 4-branch forward, backward, decode (was 3-branch, missing magnitude) - branch_actions buffer B*3 → B*4 Action diversity infrastructure: - 3-cycle counterfactual: dir mirror / mag relabel / order-type cost delta (50× amplified) - Per-branch entropy: order (d==2) gets 3× boost in C51 + MSE grad kernels - Position histogram expanded 6→12 bins (all 4 branches get entropy bonus) - Boltzmann temperature floor raised to 0.5 for order + urgency branches - Smoke test epochs 3→10 for diversity verification Note: experience collector still needs bottleneck forward path — currently reads trainer's bottleneck-layout weights with state_dim K, producing misaligned Q-values. Next: eliminate weight copy, use direct pointer views into trainer's params_buf. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -10,7 +10,7 @@
|
||||
# - q_clip ±50 (prevents logit explosion through bf16 forward)
|
||||
|
||||
[training]
|
||||
epochs = 3
|
||||
epochs = 10
|
||||
batch_size = 64
|
||||
learning_rate = 2e-6
|
||||
gamma = 0.95
|
||||
|
||||
@@ -51,7 +51,10 @@ extern "C" __global__ void c51_grad_kernel(
|
||||
* atom distribution collapse. Standardized advantages prevent Q-value scale
|
||||
* collapse; this prevents atom sharpness collapse. Two-pronged defense. */
|
||||
if (entropy_coeff > 0.0f) {
|
||||
float ent_scale = (d == 1) ? 5.0f : 1.0f;
|
||||
float ent_scale;
|
||||
if (d == 1) ent_scale = 5.0f; /* magnitude: existing */
|
||||
else if (d == 2) ent_scale = 3.0f; /* order: prevent order type collapse */
|
||||
else ent_scale = 1.0f; /* direction, urgency: baseline */
|
||||
float lp_clamped = fmaxf(lp, -10.0f);
|
||||
d_combined += ent_scale * entropy_coeff * (1.0f + lp_clamped);
|
||||
}
|
||||
|
||||
@@ -925,7 +925,12 @@ extern "C" __global__ void experience_action_select(
|
||||
q_max_ord = fmaxf(q_max_ord, q_b2[a]);
|
||||
q_min_ord = fminf(q_min_ord, q_b2[a]);
|
||||
}
|
||||
float tau_ord = fmaxf(q_max_ord - q_min_ord, 0.01f);
|
||||
/* Higher temperature floor (0.5 vs 0.01) for order: Q-values for different
|
||||
* order types start nearly identical (same reward signal). The low 0.01 floor
|
||||
* amplifies initialization noise 100×, locking in a single order type before
|
||||
* learning can differentiate them. 0.5 gives meaningful Boltzmann stochasticity
|
||||
* until the order branch develops genuine Q-gaps from counterfactual training. */
|
||||
float tau_ord = fmaxf(q_max_ord - q_min_ord, 0.5f);
|
||||
float sum_e = 0.0f;
|
||||
float exps_ord[3];
|
||||
for (int a = 0; a < b2_size; a++) {
|
||||
@@ -953,7 +958,7 @@ extern "C" __global__ void experience_action_select(
|
||||
q_max_urg = fmaxf(q_max_urg, q_b3[a]);
|
||||
q_min_urg = fminf(q_min_urg, q_b3[a]);
|
||||
}
|
||||
float tau_urg = fmaxf(q_max_urg - q_min_urg, 0.01f);
|
||||
float tau_urg = fmaxf(q_max_urg - q_min_urg, 0.5f);
|
||||
float sum_e = 0.0f;
|
||||
float exps_urg[3];
|
||||
for (int a = 0; a < b3_size; a++) {
|
||||
@@ -1661,48 +1666,44 @@ extern "C" __global__ void experience_env_step(
|
||||
reward = fmaxf(-10.0f, -5.0f - 5.0f * risk_taken);
|
||||
}
|
||||
|
||||
/* ── #19 Position entropy: track direction(3) + magnitude(3) diversity ── */
|
||||
/* ── #19 Position entropy: track all 4 branches (3 bins each = 12 total) ── */
|
||||
if (position_histogram != NULL && position_entropy_weight > 0.0f) {
|
||||
/* Histogram layout: [N, 6] — first 3 bins = direction, next 3 = magnitude */
|
||||
if (dir_idx >= 0 && dir_idx < 3) {
|
||||
position_histogram[(long long)i * 6 + dir_idx] += 1.0f;
|
||||
}
|
||||
if (mag_idx >= 0 && mag_idx < 3) {
|
||||
position_histogram[(long long)i * 6 + 3 + mag_idx] += 1.0f;
|
||||
}
|
||||
/* Histogram layout: [N, 12] — dir[0:3], mag[3:6], order[6:9], urgency[9:12] */
|
||||
int order_idx = decode_order_4b(action_idx, b2_size, b3_size);
|
||||
int urgency_idx = decode_urgency_4b(action_idx, b3_size);
|
||||
if (dir_idx >= 0 && dir_idx < 3)
|
||||
position_histogram[(long long)i * 12 + dir_idx] += 1.0f;
|
||||
if (mag_idx >= 0 && mag_idx < 3)
|
||||
position_histogram[(long long)i * 12 + 3 + mag_idx] += 1.0f;
|
||||
if (order_idx >= 0 && order_idx < 3)
|
||||
position_histogram[(long long)i * 12 + 6 + order_idx] += 1.0f;
|
||||
if (urgency_idx >= 0 && urgency_idx < 3)
|
||||
position_histogram[(long long)i * 12 + 9 + urgency_idx] += 1.0f;
|
||||
}
|
||||
|
||||
/* ---- Done detection ---- */
|
||||
int next_bar = bar_idx + 1;
|
||||
int done = (next_bar >= total_bars || check_capital_floor(new_portfolio_value, peak_equity)) ? 1 : 0;
|
||||
|
||||
/* #19 Position entropy bonus at episode end — direction(3) + magnitude(3) */
|
||||
/* #19 Position entropy bonus at episode end — all 4 branches */
|
||||
if (done && position_histogram != NULL && position_entropy_weight > 0.0f) {
|
||||
float* hist = position_histogram + (long long)i * 6;
|
||||
/* Compute entropy over direction bins [0:3] */
|
||||
float dir_total = 0.0f;
|
||||
for (int b = 0; b < 3; b++) dir_total += hist[b];
|
||||
float dir_entropy = 0.0f;
|
||||
if (dir_total > 1.0f) {
|
||||
for (int b = 0; b < 3; b++) {
|
||||
float p = hist[b] / dir_total;
|
||||
if (p > 1e-6f) dir_entropy -= p * logf(p);
|
||||
float* hist = position_histogram + (long long)i * 12;
|
||||
float combined_entropy = 0.0f;
|
||||
for (int branch = 0; branch < 4; branch++) {
|
||||
float* bh = hist + branch * 3;
|
||||
float total = 0.0f;
|
||||
for (int b = 0; b < 3; b++) total += bh[b];
|
||||
float ent = 0.0f;
|
||||
if (total > 1.0f) {
|
||||
for (int b = 0; b < 3; b++) {
|
||||
float p = bh[b] / total;
|
||||
if (p > 1e-6f) ent -= p * logf(p);
|
||||
}
|
||||
ent /= 1.0986f; /* log(3) normalization → [0, 1] */
|
||||
}
|
||||
dir_entropy /= 1.0986f; /* Normalize by log(3) = 1.0986 → [0, 1] */
|
||||
combined_entropy += ent;
|
||||
}
|
||||
/* Compute entropy over magnitude bins [3:6] */
|
||||
float mag_total = 0.0f;
|
||||
for (int b = 3; b < 6; b++) mag_total += hist[b];
|
||||
float mag_entropy = 0.0f;
|
||||
if (mag_total > 1.0f) {
|
||||
for (int b = 3; b < 6; b++) {
|
||||
float p = hist[b] / mag_total;
|
||||
if (p > 1e-6f) mag_entropy -= p * logf(p);
|
||||
}
|
||||
mag_entropy /= 1.0986f; /* Normalize by log(3) → [0, 1] */
|
||||
}
|
||||
/* Average of direction and magnitude entropy */
|
||||
float combined_entropy = 0.5f * (dir_entropy + mag_entropy);
|
||||
combined_entropy *= 0.25f; /* average across 4 branches */
|
||||
reward += position_entropy_weight * combined_entropy;
|
||||
}
|
||||
|
||||
@@ -1796,15 +1797,16 @@ extern "C" __global__ void experience_env_step(
|
||||
int cf_action;
|
||||
float cf_reward;
|
||||
|
||||
if ((current_t & 1) == 0) {
|
||||
/* ── Even step: directional mirror (original logic) ── */
|
||||
int cf_cycle = current_t % 3;
|
||||
if (cf_cycle == 0) {
|
||||
/* ── Step 0: directional mirror (original logic) ── */
|
||||
int cf_dir = (b0_size - 1) - dir_idx;
|
||||
cf_action = cf_dir * b1_size * b2_size * b3_size
|
||||
+ mag_idx * b2_size * b3_size
|
||||
+ orig_order * b3_size + orig_urgency;
|
||||
cf_reward = -reward;
|
||||
} else {
|
||||
/* ── Odd step: hindsight magnitude relabeling ──
|
||||
} else if (cf_cycle == 1) {
|
||||
/* ── Step 1: hindsight magnitude relabeling ──
|
||||
* Pick the best alternative magnitude:
|
||||
* If Full(2), try Half(1). Otherwise try Full(2).
|
||||
* Scale reward linearly by (alt_mag_frac / actual_mag_frac).
|
||||
@@ -1818,7 +1820,7 @@ extern "C" __global__ void experience_env_step(
|
||||
/* Cycle: smaller, larger, opposite. All 3 magnitudes get gradient.
|
||||
* Was: only Half or Full — Small NEVER appeared as counterfactual. */
|
||||
int alt_mag;
|
||||
int cycle = (current_t / 2) % 3;
|
||||
int cycle = (current_t / 3) % 3;
|
||||
if (cycle == 0) alt_mag = (mag_idx > 0) ? mag_idx - 1 : 2;
|
||||
else if (cycle == 1) alt_mag = (mag_idx < 2) ? mag_idx + 1 : 0;
|
||||
else alt_mag = 2 - mag_idx;
|
||||
@@ -1837,6 +1839,36 @@ extern "C" __global__ void experience_env_step(
|
||||
+ orig_order * b3_size + orig_urgency;
|
||||
cf_reward = -reward;
|
||||
}
|
||||
} else {
|
||||
/* ── Step 2: Order-type counterfactual ──
|
||||
* Cycle through alternative order types. Compute execution cost
|
||||
* difference and amplify to match directional signal magnitude. */
|
||||
int alt_order = (orig_order + 1 + ((current_t / 3) & 1)) % b2_size;
|
||||
|
||||
if (fabsf(position) > 0.001f) {
|
||||
float taken_cost = compute_tx_cost(
|
||||
position - pre_trade_position, raw_close, tx_cost_multiplier,
|
||||
spread_cost, max_position, orig_order, spread_scale);
|
||||
float alt_cost = compute_tx_cost(
|
||||
position - pre_trade_position, raw_close, tx_cost_multiplier,
|
||||
spread_cost, max_position, alt_order, spread_scale);
|
||||
float cost_delta = (taken_cost - alt_cost)
|
||||
/ fmaxf(prev_equity, 1.0f);
|
||||
|
||||
cf_action = dir_idx * b1_size * b2_size * b3_size
|
||||
+ mag_idx * b2_size * b3_size
|
||||
+ alt_order * b3_size + orig_urgency;
|
||||
/* Amplify 50×: raw cost delta ~0.0001-0.0005, directional ~0.01-0.1 */
|
||||
cf_reward = reward + 50.0f * cost_delta;
|
||||
cf_reward = fminf(fmaxf(cf_reward, -10.0f), 10.0f);
|
||||
} else {
|
||||
/* Flat position: order irrelevant, directional mirror fallback */
|
||||
int cf_dir = (b0_size - 1) - dir_idx;
|
||||
cf_action = cf_dir * b1_size * b2_size * b3_size
|
||||
+ mag_idx * b2_size * b3_size
|
||||
+ orig_order * b3_size + orig_urgency;
|
||||
cf_reward = -reward;
|
||||
}
|
||||
}
|
||||
|
||||
if (isnan(cf_reward) || isinf(cf_reward)) cf_reward = 0.0f;
|
||||
|
||||
@@ -1045,8 +1045,8 @@ impl GpuExperienceCollector {
|
||||
"GPU experience collector initialized (cuBLAS + 3 focused kernels)"
|
||||
);
|
||||
|
||||
// #19 Position entropy histogram [alloc_episodes * 6] — direction(3) + magnitude(3)
|
||||
let position_histogram = stream.alloc_zeros::<f32>(alloc_episodes * 6)
|
||||
// #19 Position entropy histogram [alloc_episodes * 12] — dir(3)+mag(3)+order(3)+urgency(3)
|
||||
let position_histogram = stream.alloc_zeros::<f32>(alloc_episodes * 12)
|
||||
.map_err(|e| MLError::ModelError(format!("alloc position_histogram: {e}")))?;
|
||||
|
||||
// #30 F32 forward pass buffers for full-precision experience collection
|
||||
|
||||
@@ -178,7 +178,7 @@ pub struct GpuIqnHead {
|
||||
/// Precomputed cosine features [N, embed_dim] — cos(π·(d+1)·τ_i)
|
||||
/// Fixed at construction (τ are midpoints). Eliminates 16.7M cosf() per step.
|
||||
cos_features: CudaSlice<f32>,
|
||||
/// Branch actions decoded from flat actions [B, 3]
|
||||
/// Branch actions decoded from flat actions [B, 4]
|
||||
branch_actions: CudaSlice<i32>,
|
||||
/// Target h_s2 computed from next_states + target trunk weights [B, H]
|
||||
target_h_s2: CudaSlice<f32>,
|
||||
@@ -282,7 +282,7 @@ impl GpuIqnHead {
|
||||
MLError::ModelError(format!("IQN htod target_taus ({} f32): {e}", b * n))
|
||||
})?;
|
||||
}
|
||||
let branch_actions = stream.alloc_zeros::<i32>(b * 3).map_err(|e| {
|
||||
let branch_actions = stream.alloc_zeros::<i32>(b * 4).map_err(|e| {
|
||||
MLError::ModelError(format!("IQN alloc branch_actions: {e}"))
|
||||
})?;
|
||||
let target_h_s2 = alloc_f32(&stream, b * h, "iqn_target_h_s2")?;
|
||||
@@ -306,9 +306,9 @@ impl GpuIqnHead {
|
||||
// IQR buffer: per-action interquartile range for exploration bonus
|
||||
let iqr_buf = alloc_f32(&stream, tba, "iqn_iqr")?;
|
||||
|
||||
let vram_bytes = (total_params * 6 + b * n * 2 + b * 3
|
||||
let vram_bytes = (total_params * 6 + b * n * 2 + b * 4
|
||||
+ b * h + b * n * h * 2 + b * n * tba * 2
|
||||
+ b * 3 + 2) * 4;
|
||||
+ b * 4 + 2) * 4;
|
||||
|
||||
info!(
|
||||
hidden_dim = h,
|
||||
@@ -395,11 +395,15 @@ impl GpuIqnHead {
|
||||
let shared_h1_i32 = self.config.shared_h1 as i32;
|
||||
let hidden_dim_i32 = self.config.hidden_dim as i32;
|
||||
let embed_dim_i32 = self.config.embed_dim as i32;
|
||||
let b0_i32 = self.config.branch_0_size as i32;
|
||||
let b1_i32 = self.config.branch_1_size as i32;
|
||||
let b2_i32 = self.config.branch_2_size as i32;
|
||||
let b3_i32 = self.config.branch_3_size as i32;
|
||||
|
||||
// Layer 3a: τ values are fixed midpoints, pre-computed in constructor.
|
||||
// No per-step sampling → deterministic → CUDA Graph compatible.
|
||||
|
||||
// 2. GPU decode: flat actions [B] → branch actions [B, 3]
|
||||
// 2. GPU decode: flat actions [B] → branch actions [B, 4]
|
||||
let decode_blocks = (b + 255) / 256;
|
||||
let decode_config = LaunchConfig {
|
||||
grid_dim: (decode_blocks as u32, 1, 1),
|
||||
@@ -408,7 +412,7 @@ impl GpuIqnHead {
|
||||
};
|
||||
let batch_size_i32 = b as i32;
|
||||
|
||||
// Safety: dqn_actions_buf[B] i32, branch_actions[B*3] i32, same context.
|
||||
// Safety: dqn_actions_buf[B] i32, branch_actions[B*4] i32, same context.
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.decode_actions_kernel)
|
||||
@@ -418,6 +422,10 @@ impl GpuIqnHead {
|
||||
.arg(&shared_h1_i32)
|
||||
.arg(&hidden_dim_i32)
|
||||
.arg(&embed_dim_i32)
|
||||
.arg(&b0_i32)
|
||||
.arg(&b1_i32)
|
||||
.arg(&b2_i32)
|
||||
.arg(&b3_i32)
|
||||
.launch(decode_config)
|
||||
.map_err(|e| MLError::ModelError(format!("IQN decode_actions kernel: {e}")))?;
|
||||
}
|
||||
@@ -461,6 +469,10 @@ impl GpuIqnHead {
|
||||
let shared_h1_i32 = self.config.shared_h1 as i32;
|
||||
let hidden_dim_i32 = self.config.hidden_dim as i32;
|
||||
let embed_dim_i32 = self.config.embed_dim as i32;
|
||||
let b0_i32 = self.config.branch_0_size as i32;
|
||||
let b1_i32 = self.config.branch_1_size as i32;
|
||||
let b2_i32 = self.config.branch_2_size as i32;
|
||||
let b3_i32 = self.config.branch_3_size as i32;
|
||||
|
||||
if let Some(cached_ptr) = self.cached_target_h_s2_ptr.take() {
|
||||
// Reuse pre-computed target h_s2 from DQN Pass 2 — skip trunk forward.
|
||||
@@ -496,6 +508,10 @@ impl GpuIqnHead {
|
||||
.arg(&shared_h1_i32)
|
||||
.arg(&hidden_dim_i32)
|
||||
.arg(&embed_dim_i32)
|
||||
.arg(&b0_i32)
|
||||
.arg(&b1_i32)
|
||||
.arg(&b2_i32)
|
||||
.arg(&b3_i32)
|
||||
.launch(trunk_config)
|
||||
.map_err(|e| MLError::ModelError(format!("IQN trunk forward kernel: {e}")))?;
|
||||
}
|
||||
@@ -549,6 +565,10 @@ impl GpuIqnHead {
|
||||
.arg(&shared_h1_i32)
|
||||
.arg(&hidden_dim_i32)
|
||||
.arg(&embed_dim_i32)
|
||||
.arg(&b0_i32)
|
||||
.arg(&b1_i32)
|
||||
.arg(&b2_i32)
|
||||
.arg(&b3_i32)
|
||||
.launch(fwd_config)
|
||||
.map_err(|e| MLError::ModelError(format!("IQN forward+loss kernel: {e}")))?;
|
||||
}
|
||||
@@ -580,6 +600,10 @@ impl GpuIqnHead {
|
||||
.arg(&shared_h1_i32)
|
||||
.arg(&hidden_dim_i32)
|
||||
.arg(&embed_dim_i32)
|
||||
.arg(&b0_i32)
|
||||
.arg(&b1_i32)
|
||||
.arg(&b2_i32)
|
||||
.arg(&b3_i32)
|
||||
.launch(bwd_config)
|
||||
.map_err(|e| MLError::ModelError(format!("IQN backward kernel: {e}")))?;
|
||||
}
|
||||
@@ -603,6 +627,10 @@ impl GpuIqnHead {
|
||||
.arg(&shared_h1_i32)
|
||||
.arg(&hidden_dim_i32)
|
||||
.arg(&embed_dim_i32)
|
||||
.arg(&b0_i32)
|
||||
.arg(&b1_i32)
|
||||
.arg(&b2_i32)
|
||||
.arg(&b3_i32)
|
||||
.launch(norm_config)
|
||||
.map_err(|e| MLError::ModelError(format!("IQN grad_norm kernel: {e}")))?;
|
||||
}
|
||||
@@ -652,6 +680,10 @@ impl GpuIqnHead {
|
||||
.arg(&shared_h1_i32)
|
||||
.arg(&hidden_dim_i32)
|
||||
.arg(&embed_dim_i32)
|
||||
.arg(&b0_i32)
|
||||
.arg(&b1_i32)
|
||||
.arg(&b2_i32)
|
||||
.arg(&b3_i32)
|
||||
.launch(adam_config)
|
||||
.map_err(|e| MLError::ModelError(format!("IQN adam kernel: {e}")))?;
|
||||
}
|
||||
@@ -687,6 +719,10 @@ impl GpuIqnHead {
|
||||
let shared_h1_i32 = self.config.shared_h1 as i32;
|
||||
let hidden_dim_i32 = self.config.hidden_dim as i32;
|
||||
let embed_dim_i32 = self.config.embed_dim as i32;
|
||||
let b0_i32 = self.config.branch_0_size as i32;
|
||||
let b1_i32 = self.config.branch_1_size as i32;
|
||||
let b2_i32 = self.config.branch_2_size as i32;
|
||||
let b3_i32 = self.config.branch_3_size as i32;
|
||||
let tau_ptr = self.tau_buf.raw_ptr();
|
||||
|
||||
// Safety: target_params and online_params have total_params elements each.
|
||||
@@ -700,6 +736,10 @@ impl GpuIqnHead {
|
||||
.arg(&shared_h1_i32)
|
||||
.arg(&hidden_dim_i32)
|
||||
.arg(&embed_dim_i32)
|
||||
.arg(&b0_i32)
|
||||
.arg(&b1_i32)
|
||||
.arg(&b2_i32)
|
||||
.arg(&b3_i32)
|
||||
.launch(config)
|
||||
.map_err(|e| MLError::ModelError(format!("IQN EMA kernel: {e}")))?;
|
||||
}
|
||||
@@ -852,6 +892,10 @@ impl GpuIqnHead {
|
||||
block_dim: (256, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
let b0_i32 = self.config.branch_0_size as i32;
|
||||
let b1_i32 = self.config.branch_1_size as i32;
|
||||
let b2_i32 = self.config.branch_2_size as i32;
|
||||
let b3_i32 = self.config.branch_3_size as i32;
|
||||
let rng_step = self.rng_step as u32;
|
||||
self.rng_step += 1;
|
||||
unsafe {
|
||||
@@ -863,6 +907,10 @@ impl GpuIqnHead {
|
||||
.arg(&shared_h1_i32)
|
||||
.arg(&hidden_dim_i32)
|
||||
.arg(&embed_dim_i32)
|
||||
.arg(&b0_i32)
|
||||
.arg(&b1_i32)
|
||||
.arg(&b2_i32)
|
||||
.arg(&b3_i32)
|
||||
.launch(tau_config)
|
||||
.map_err(|e| MLError::ModelError(format!("IQN CVaR sample_taus: {e}")))?;
|
||||
}
|
||||
@@ -887,6 +935,10 @@ impl GpuIqnHead {
|
||||
.arg(&shared_h1_i32)
|
||||
.arg(&hidden_dim_i32)
|
||||
.arg(&embed_dim_i32)
|
||||
.arg(&b0_i32)
|
||||
.arg(&b1_i32)
|
||||
.arg(&b2_i32)
|
||||
.arg(&b3_i32)
|
||||
.launch(fwd_config)
|
||||
.map_err(|e| MLError::ModelError(format!("IQN CVaR forward: {e}")))?;
|
||||
}
|
||||
|
||||
@@ -34,9 +34,7 @@
|
||||
* IQN_EMBED_DIM -- cosine embedding dimension (default 64)
|
||||
* IQN_NUM_QUANTILES -- τ samples per forward pass (default 32)
|
||||
* IQN_KAPPA -- Huber threshold for quantile loss (default 1.0)
|
||||
* BRANCH_0_SIZE -- exposure actions (default 5)
|
||||
* BRANCH_1_SIZE -- order type actions (default 3)
|
||||
* BRANCH_2_SIZE -- urgency actions (default 3)
|
||||
* Branch sizes (b0, b1, b2, b3) are runtime kernel parameters.
|
||||
*/
|
||||
|
||||
/* ── Compile-time defaults ──────────────────────────────────────────── */
|
||||
@@ -59,50 +57,17 @@
|
||||
#ifndef IQN_KAPPA
|
||||
#define IQN_KAPPA 1.0f
|
||||
#endif
|
||||
#ifndef BRANCH_0_SIZE
|
||||
#define BRANCH_0_SIZE 5
|
||||
#endif
|
||||
#ifndef BRANCH_1_SIZE
|
||||
#define BRANCH_1_SIZE 3
|
||||
#endif
|
||||
#ifndef BRANCH_2_SIZE
|
||||
#define BRANCH_2_SIZE 3
|
||||
#endif
|
||||
|
||||
#define TOTAL_BRANCH_ACTIONS (BRANCH_0_SIZE + BRANCH_1_SIZE + BRANCH_2_SIZE)
|
||||
/* Branch sizes (b0_size, b1_size, b2_size, b3_size) are runtime kernel parameters.
|
||||
* No compile-time BRANCH_*_SIZE defines — all offsets computed at runtime. */
|
||||
|
||||
/* ── Weight layout (flat f32 buffer) ─────────────────────────────────── */
|
||||
/* W_embed [hidden_dim, IQN_EMBED_DIM], b_embed [hidden_dim]
|
||||
* W_b0 [BRANCH_0_SIZE, hidden_dim], b_b0 [BRANCH_0_SIZE]
|
||||
* W_b1 [BRANCH_1_SIZE, hidden_dim], b_b1 [BRANCH_1_SIZE]
|
||||
* W_b2 [BRANCH_2_SIZE, hidden_dim], b_b2 [BRANCH_2_SIZE]
|
||||
/* W_embed [hidden_dim, embed_dim], b_embed [hidden_dim]
|
||||
* W_b0 [b0_size, hidden_dim], b_b0 [b0_size]
|
||||
* W_b1 [b1_size, hidden_dim], b_b1 [b1_size]
|
||||
* W_b2 [b2_size, hidden_dim], b_b2 [b2_size]
|
||||
* W_b3 [b3_size, hidden_dim], b_b3 [b3_size]
|
||||
*
|
||||
* NOTE: IQN_OFF_* macros below use compile-time IQN_HIDDEN. They are kept
|
||||
* only for documentation. Kernels compute offsets from runtime hidden_dim. */
|
||||
|
||||
#define IQN_W_EMBED_SIZE (IQN_HIDDEN * IQN_EMBED_DIM)
|
||||
#define IQN_B_EMBED_SIZE (IQN_HIDDEN)
|
||||
#define IQN_W_B0_SIZE (BRANCH_0_SIZE * IQN_HIDDEN)
|
||||
#define IQN_B_B0_SIZE (BRANCH_0_SIZE)
|
||||
#define IQN_W_B1_SIZE (BRANCH_1_SIZE * IQN_HIDDEN)
|
||||
#define IQN_B_B1_SIZE (BRANCH_1_SIZE)
|
||||
#define IQN_W_B2_SIZE (BRANCH_2_SIZE * IQN_HIDDEN)
|
||||
#define IQN_B_B2_SIZE (BRANCH_2_SIZE)
|
||||
|
||||
#define IQN_TOTAL_PARAMS (IQN_W_EMBED_SIZE + IQN_B_EMBED_SIZE \
|
||||
+ IQN_W_B0_SIZE + IQN_B_B0_SIZE \
|
||||
+ IQN_W_B1_SIZE + IQN_B_B1_SIZE \
|
||||
+ IQN_W_B2_SIZE + IQN_B_B2_SIZE)
|
||||
|
||||
/* Offsets into the flat parameter buffer (compile-time, for documentation only) */
|
||||
#define IQN_OFF_W_EMBED 0
|
||||
#define IQN_OFF_B_EMBED (IQN_OFF_W_EMBED + IQN_W_EMBED_SIZE)
|
||||
#define IQN_OFF_W_B0 (IQN_OFF_B_EMBED + IQN_B_EMBED_SIZE)
|
||||
#define IQN_OFF_B_B0 (IQN_OFF_W_B0 + IQN_W_B0_SIZE)
|
||||
#define IQN_OFF_W_B1 (IQN_OFF_B_B0 + IQN_B_B0_SIZE)
|
||||
#define IQN_OFF_B_B1 (IQN_OFF_W_B1 + IQN_W_B1_SIZE)
|
||||
#define IQN_OFF_W_B2 (IQN_OFF_B_B1 + IQN_B_B1_SIZE)
|
||||
#define IQN_OFF_B_B2 (IQN_OFF_W_B2 + IQN_W_B2_SIZE)
|
||||
* Offsets computed at runtime by iqn_compute_offsets(). */
|
||||
|
||||
/* Block size for IQN forward/backward/trunk kernels (8 warps = 256 threads). */
|
||||
#define IQN_BLOCK_SIZE 256
|
||||
@@ -185,17 +150,25 @@ __device__ __forceinline__ float quantile_huber_grad(float tau, float delta, flo
|
||||
return -weight * huber_grad;
|
||||
}
|
||||
|
||||
/** Compute runtime weight offsets from hidden_dim and embed_dim. Returns offsets in an array:
|
||||
* [0]=W_EMBED, [1]=B_EMBED, [2]=W_B0, [3]=B_B0, [4]=W_B1, [5]=B_B1, [6]=W_B2, [7]=B_B2 */
|
||||
__device__ __forceinline__ void iqn_compute_offsets(int hidden_dim, int embed_dim, int offsets[8]) {
|
||||
/** Compute runtime weight offsets from hidden_dim, embed_dim, and 4 branch sizes.
|
||||
* Returns 10 offsets:
|
||||
* [0]=W_EMBED, [1]=B_EMBED, [2]=W_B0, [3]=B_B0, [4]=W_B1, [5]=B_B1,
|
||||
* [6]=W_B2, [7]=B_B2, [8]=W_B3, [9]=B_B3 */
|
||||
__device__ __forceinline__ void iqn_compute_offsets(
|
||||
int hidden_dim, int embed_dim,
|
||||
int b0, int b1, int b2, int b3,
|
||||
int offsets[10])
|
||||
{
|
||||
offsets[0] = 0; /* W_EMBED */
|
||||
offsets[1] = offsets[0] + hidden_dim * embed_dim; /* B_EMBED */
|
||||
offsets[2] = offsets[1] + hidden_dim; /* W_B0 */
|
||||
offsets[3] = offsets[2] + BRANCH_0_SIZE * hidden_dim; /* B_B0 */
|
||||
offsets[4] = offsets[3] + BRANCH_0_SIZE; /* W_B1 */
|
||||
offsets[5] = offsets[4] + BRANCH_1_SIZE * hidden_dim; /* B_B1 */
|
||||
offsets[6] = offsets[5] + BRANCH_1_SIZE; /* W_B2 */
|
||||
offsets[7] = offsets[6] + BRANCH_2_SIZE * hidden_dim; /* B_B2 */
|
||||
offsets[3] = offsets[2] + b0 * hidden_dim; /* B_B0 */
|
||||
offsets[4] = offsets[3] + b0; /* W_B1 */
|
||||
offsets[5] = offsets[4] + b1 * hidden_dim; /* B_B1 */
|
||||
offsets[6] = offsets[5] + b1; /* W_B2 */
|
||||
offsets[7] = offsets[6] + b2 * hidden_dim; /* B_B2 */
|
||||
offsets[8] = offsets[7] + b2; /* W_B3 */
|
||||
offsets[9] = offsets[8] + b3 * hidden_dim; /* B_B3 */
|
||||
}
|
||||
|
||||
/* ── Forward + Loss Kernel ──────────────────────────────────────────── */
|
||||
@@ -208,7 +181,7 @@ __device__ __forceinline__ void iqn_compute_offsets(int hidden_dim, int embed_di
|
||||
* τ values are pre-sampled and uploaded via the `taus` buffer (f32).
|
||||
*
|
||||
* One warp (32 threads) per sample. Quantiles are processed sequentially
|
||||
* to conserve registers — each τ produces TOTAL_BRANCH_ACTIONS values.
|
||||
* to conserve registers — each τ produces tba (= b0+b1+b2+b3) values.
|
||||
*
|
||||
* Saves per-quantile intermediate activations for backward pass.
|
||||
*/
|
||||
@@ -220,7 +193,7 @@ void iqn_forward_loss_kernel(
|
||||
/* Inputs (f32) */
|
||||
const float* __restrict__ taus, /* [B, N] online τ samples ∈ (0,1) */
|
||||
const float* __restrict__ target_taus, /* [B, N] target τ samples */
|
||||
const int* __restrict__ actions, /* [B, 3] branch actions (exposure, order, urgency) */
|
||||
const int* __restrict__ actions, /* [B, 4] branch actions (direction, magnitude, order, urgency) */
|
||||
const float* __restrict__ rewards, /* [B] */
|
||||
const float* __restrict__ dones, /* [B] (0.0/1.0) */
|
||||
float gamma,
|
||||
@@ -235,12 +208,16 @@ void iqn_forward_loss_kernel(
|
||||
/* Activation saves for backward (f32) */
|
||||
float* __restrict__ save_embed, /* [B, N, hidden_dim] post-ReLU embeddings */
|
||||
float* __restrict__ save_combined, /* [B, N, hidden_dim] h_s2 ⊙ embed */
|
||||
float* __restrict__ save_q_online, /* [B, N, TOTAL_BRANCH_ACTIONS] quantile Q-values */
|
||||
float* __restrict__ save_q_target, /* [B, N, TOTAL_BRANCH_ACTIONS] target Q-values */
|
||||
float* __restrict__ save_q_online, /* [B, N, tba] quantile Q-values */
|
||||
float* __restrict__ save_q_target, /* [B, N, tba] target Q-values */
|
||||
int batch_size,
|
||||
int shared_h1, /* runtime: first hidden layer width (unused here, for consistency) */
|
||||
int hidden_dim, /* runtime: IQN hidden dim = SHARED_H2 */
|
||||
int embed_dim /* runtime: cosine embedding dimension (was IQN_EMBED_DIM) */
|
||||
int embed_dim, /* runtime: cosine embedding dimension (was IQN_EMBED_DIM) */
|
||||
int b0_size, /* runtime: branch 0 (direction) action count */
|
||||
int b1_size, /* runtime: branch 1 (magnitude) action count */
|
||||
int b2_size, /* runtime: branch 2 (order) action count */
|
||||
int b3_size /* runtime: branch 3 (urgency) action count */
|
||||
)
|
||||
{
|
||||
int sample = blockIdx.x;
|
||||
@@ -251,8 +228,9 @@ void iqn_forward_loss_kernel(
|
||||
__shared__ float shmem_reduce[IQN_BLOCK_SIZE / 32];
|
||||
|
||||
/* ── Compute runtime weight offsets ── */
|
||||
int off[8];
|
||||
iqn_compute_offsets(hidden_dim, embed_dim, off);
|
||||
int tba = b0_size + b1_size + b2_size + b3_size;
|
||||
int off[10];
|
||||
iqn_compute_offsets(hidden_dim, embed_dim, b0_size, b1_size, b2_size, b3_size, off);
|
||||
|
||||
/* ── Pointers into this sample's data ── */
|
||||
const float* my_h_s2_bf16 = h_s2 + sample * hidden_dim;
|
||||
@@ -260,9 +238,10 @@ void iqn_forward_loss_kernel(
|
||||
const float* my_taus = taus + sample * IQN_NUM_QUANTILES;
|
||||
const float* my_target_taus = target_taus + sample * IQN_NUM_QUANTILES;
|
||||
|
||||
int a0 = actions[sample * 3 + 0]; /* exposure action */
|
||||
int a1 = actions[sample * 3 + 1]; /* order action */
|
||||
int a2 = actions[sample * 3 + 2]; /* urgency action */
|
||||
int a0 = actions[sample * 4 + 0]; /* direction action */
|
||||
int a1 = actions[sample * 4 + 1]; /* magnitude action */
|
||||
int a2 = actions[sample * 4 + 2]; /* order action */
|
||||
int a3 = actions[sample * 4 + 3]; /* urgency action */
|
||||
float reward = rewards[sample];
|
||||
float done = dones[sample];
|
||||
|
||||
@@ -275,6 +254,8 @@ void iqn_forward_loss_kernel(
|
||||
const float* b_b1 = online_params + off[5];
|
||||
const float* w_b2 = online_params + off[6];
|
||||
const float* b_b2 = online_params + off[7];
|
||||
const float* w_b3 = online_params + off[8];
|
||||
const float* b_b3 = online_params + off[9];
|
||||
|
||||
const float* tw_embed = target_params + off[0];
|
||||
const float* tb_embed = target_params + off[1];
|
||||
@@ -284,6 +265,8 @@ void iqn_forward_loss_kernel(
|
||||
const float* tb_b1 = target_params + off[5];
|
||||
const float* tw_b2 = target_params + off[6];
|
||||
const float* tb_b2 = target_params + off[7];
|
||||
const float* tw_b3 = target_params + off[8];
|
||||
const float* tb_b3 = target_params + off[9];
|
||||
|
||||
/* ── Load h_s2 into distributed registers (bf16 → f32 at boundary) ── */
|
||||
float h_dist[IQN_DIST_MAX];
|
||||
@@ -328,11 +311,11 @@ void iqn_forward_loss_kernel(
|
||||
}
|
||||
|
||||
/* 4. Per-branch output: q_d = W_bd × combined + b_bd */
|
||||
int q_save_offset = (sample * IQN_NUM_QUANTILES + t) * TOTAL_BRANCH_ACTIONS;
|
||||
int q_save_offset = (sample * IQN_NUM_QUANTILES + t) * tba;
|
||||
|
||||
/* Branch 0 (exposure) */
|
||||
/* Branch 0 (direction) */
|
||||
float q0_taken = 0.0f;
|
||||
for (int a = 0; a < BRANCH_0_SIZE; a++) {
|
||||
for (int a = 0; a < b0_size; a++) {
|
||||
float partial = 0.0f;
|
||||
const float* w_row = w_b0 + a * hidden_dim;
|
||||
for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE)
|
||||
@@ -348,16 +331,16 @@ void iqn_forward_loss_kernel(
|
||||
__syncthreads();
|
||||
q0_taken = shmem_reduce[0];
|
||||
|
||||
/* Branch 1 (order) */
|
||||
/* Branch 1 (magnitude) */
|
||||
float q1_taken = 0.0f;
|
||||
for (int a = 0; a < BRANCH_1_SIZE; a++) {
|
||||
for (int a = 0; a < b1_size; a++) {
|
||||
float partial = 0.0f;
|
||||
const float* w_row = w_b1 + a * hidden_dim;
|
||||
for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE)
|
||||
partial += w_row[h] * comb_dist[h / IQN_BLOCK_SIZE];
|
||||
float acc = b_b1[a] + iqn_block_sum(partial, shmem_reduce);
|
||||
if (tid == 0) {
|
||||
save_q_online[q_save_offset + BRANCH_0_SIZE + a] = acc;
|
||||
save_q_online[q_save_offset + b0_size + a] = acc;
|
||||
if (a == a1) q1_taken = acc;
|
||||
}
|
||||
}
|
||||
@@ -365,16 +348,16 @@ void iqn_forward_loss_kernel(
|
||||
__syncthreads();
|
||||
q1_taken = shmem_reduce[0];
|
||||
|
||||
/* Branch 2 (urgency) */
|
||||
/* Branch 2 (order) */
|
||||
float q2_taken = 0.0f;
|
||||
for (int a = 0; a < BRANCH_2_SIZE; a++) {
|
||||
for (int a = 0; a < b2_size; a++) {
|
||||
float partial = 0.0f;
|
||||
const float* w_row = w_b2 + a * hidden_dim;
|
||||
for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE)
|
||||
partial += w_row[h] * comb_dist[h / IQN_BLOCK_SIZE];
|
||||
float acc = b_b2[a] + iqn_block_sum(partial, shmem_reduce);
|
||||
if (tid == 0) {
|
||||
save_q_online[q_save_offset + BRANCH_0_SIZE + BRANCH_1_SIZE + a] = acc;
|
||||
save_q_online[q_save_offset + b0_size + b1_size + a] = acc;
|
||||
if (a == a2) q2_taken = acc;
|
||||
}
|
||||
}
|
||||
@@ -382,8 +365,25 @@ void iqn_forward_loss_kernel(
|
||||
__syncthreads();
|
||||
q2_taken = shmem_reduce[0];
|
||||
|
||||
/* Branch 3 (urgency) */
|
||||
float q3_taken = 0.0f;
|
||||
for (int a = 0; a < b3_size; a++) {
|
||||
float partial = 0.0f;
|
||||
const float* w_row = w_b3 + a * hidden_dim;
|
||||
for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE)
|
||||
partial += w_row[h] * comb_dist[h / IQN_BLOCK_SIZE];
|
||||
float acc = b_b3[a] + iqn_block_sum(partial, shmem_reduce);
|
||||
if (tid == 0) {
|
||||
save_q_online[q_save_offset + b0_size + b1_size + b2_size + a] = acc;
|
||||
if (a == a3) q3_taken = acc;
|
||||
}
|
||||
}
|
||||
if (tid == 0) shmem_reduce[0] = q3_taken;
|
||||
__syncthreads();
|
||||
q3_taken = shmem_reduce[0];
|
||||
|
||||
/* Sum of taken-action Q-values across branches (for loss computation) */
|
||||
online_q_a[t] = q0_taken + q1_taken + q2_taken;
|
||||
online_q_a[t] = q0_taken + q1_taken + q2_taken + q3_taken;
|
||||
}
|
||||
|
||||
/* ── Process TARGET quantiles ── */
|
||||
@@ -410,7 +410,7 @@ void iqn_forward_loss_kernel(
|
||||
comb_dist[h / IQN_BLOCK_SIZE] = h_target_dist[h / IQN_BLOCK_SIZE] * embed_dist[h / IQN_BLOCK_SIZE];
|
||||
|
||||
/* Per-branch Q-values for taken actions */
|
||||
float q0, q1, q2;
|
||||
float q0, q1, q2, q3;
|
||||
|
||||
/* Branch 0 */
|
||||
{
|
||||
@@ -436,17 +436,26 @@ void iqn_forward_loss_kernel(
|
||||
partial += w_row[h] * comb_dist[h / IQN_BLOCK_SIZE];
|
||||
q2 = tb_b2[a2] + iqn_block_sum(partial, shmem_reduce);
|
||||
}
|
||||
/* Branch 3 */
|
||||
{
|
||||
float partial = 0.0f;
|
||||
const float* w_row = tw_b3 + a3 * hidden_dim;
|
||||
for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE)
|
||||
partial += w_row[h] * comb_dist[h / IQN_BLOCK_SIZE];
|
||||
q3 = tb_b3[a3] + iqn_block_sum(partial, shmem_reduce);
|
||||
}
|
||||
|
||||
/* Bellman target: r + γ(1-done) × Q_target */
|
||||
float raw_target = q0 + q1 + q2;
|
||||
float raw_target = q0 + q1 + q2 + q3;
|
||||
target_q_a[t] = reward + gamma * (1.0f - done) * raw_target;
|
||||
|
||||
/* Save target Q for backward */
|
||||
int q_save_offset = (sample * IQN_NUM_QUANTILES + t) * TOTAL_BRANCH_ACTIONS;
|
||||
int q_save_offset = (sample * IQN_NUM_QUANTILES + t) * tba;
|
||||
if (tid == 0) {
|
||||
save_q_target[q_save_offset + a0] = target_q_a[t];
|
||||
save_q_target[q_save_offset + BRANCH_0_SIZE + a1] = target_q_a[t];
|
||||
save_q_target[q_save_offset + BRANCH_0_SIZE + BRANCH_1_SIZE + a2] = target_q_a[t];
|
||||
save_q_target[q_save_offset + b0_size + a1] = target_q_a[t];
|
||||
save_q_target[q_save_offset + b0_size + b1_size + a2] = target_q_a[t];
|
||||
save_q_target[q_save_offset + b0_size + b1_size + b2_size + a3] = target_q_a[t];
|
||||
}
|
||||
}
|
||||
|
||||
@@ -491,10 +500,10 @@ void iqn_backward_kernel(
|
||||
const float* __restrict__ h_s2, /* [B, hidden_dim] (bf16 from DQN trunk) */
|
||||
const float* __restrict__ save_embed, /* [B, N, hidden_dim] */
|
||||
const float* __restrict__ save_combined, /* [B, N, hidden_dim] */
|
||||
const float* __restrict__ save_q_online, /* [B, N, TOTAL_BRANCH_ACTIONS] */
|
||||
const float* __restrict__ save_q_target, /* [B, N, TOTAL_BRANCH_ACTIONS] */
|
||||
const float* __restrict__ save_q_online, /* [B, N, tba] */
|
||||
const float* __restrict__ save_q_target, /* [B, N, tba] */
|
||||
const float* __restrict__ taus, /* [B, N] */
|
||||
const int* __restrict__ actions, /* [B, 3] */
|
||||
const int* __restrict__ actions, /* [B, 4] */
|
||||
/* Weights (f32, for backward through linear layers) */
|
||||
const float* __restrict__ online_params,
|
||||
/* Precomputed cosine features (Optimization C) */
|
||||
@@ -505,7 +514,11 @@ void iqn_backward_kernel(
|
||||
int batch_size,
|
||||
int shared_h1, /* runtime: first hidden layer width (unused here) */
|
||||
int hidden_dim, /* runtime: IQN hidden dim = SHARED_H2 */
|
||||
int embed_dim /* runtime: cosine embedding dimension (was IQN_EMBED_DIM) */
|
||||
int embed_dim, /* runtime: cosine embedding dimension (was IQN_EMBED_DIM) */
|
||||
int b0_size, /* runtime: branch 0 (direction) action count */
|
||||
int b1_size, /* runtime: branch 1 (magnitude) action count */
|
||||
int b2_size, /* runtime: branch 2 (order) action count */
|
||||
int b3_size /* runtime: branch 3 (urgency) action count */
|
||||
)
|
||||
{
|
||||
int sample = blockIdx.x;
|
||||
@@ -516,19 +529,22 @@ void iqn_backward_kernel(
|
||||
__shared__ float shmem_reduce[IQN_BLOCK_SIZE / 32];
|
||||
|
||||
/* ── Compute runtime weight offsets ── */
|
||||
int off[8];
|
||||
iqn_compute_offsets(hidden_dim, embed_dim, off);
|
||||
int tba = b0_size + b1_size + b2_size + b3_size;
|
||||
int off[10];
|
||||
iqn_compute_offsets(hidden_dim, embed_dim, b0_size, b1_size, b2_size, b3_size, off);
|
||||
|
||||
const float* my_h_s2_bf16 = h_s2 + sample * hidden_dim;
|
||||
const float* my_taus = taus + sample * IQN_NUM_QUANTILES;
|
||||
int a0 = actions[sample * 3 + 0];
|
||||
int a1 = actions[sample * 3 + 1];
|
||||
int a2 = actions[sample * 3 + 2];
|
||||
int a0 = actions[sample * 4 + 0];
|
||||
int a1 = actions[sample * 4 + 1];
|
||||
int a2 = actions[sample * 4 + 2];
|
||||
int a3 = actions[sample * 4 + 3];
|
||||
|
||||
const float* w_embed = online_params + off[0];
|
||||
const float* w_b0 = online_params + off[2];
|
||||
const float* w_b1 = online_params + off[4];
|
||||
const float* w_b2 = online_params + off[6];
|
||||
const float* w_b3 = online_params + off[8];
|
||||
|
||||
/* Load h_s2 into distributed registers (bf16 → f32 at boundary) */
|
||||
float h_dist[IQN_DIST_MAX];
|
||||
@@ -540,11 +556,12 @@ void iqn_backward_kernel(
|
||||
float online_q[IQN_NUM_QUANTILES];
|
||||
float target_q[IQN_NUM_QUANTILES];
|
||||
for (int t = 0; t < IQN_NUM_QUANTILES; t++) {
|
||||
int qoff = (sample * IQN_NUM_QUANTILES + t) * TOTAL_BRANCH_ACTIONS;
|
||||
int qoff = (sample * IQN_NUM_QUANTILES + t) * tba;
|
||||
if (tid == 0) {
|
||||
online_q[t] = save_q_online[qoff + a0]
|
||||
+ save_q_online[qoff + BRANCH_0_SIZE + a1]
|
||||
+ save_q_online[qoff + BRANCH_0_SIZE + BRANCH_1_SIZE + a2];
|
||||
+ save_q_online[qoff + b0_size + a1]
|
||||
+ save_q_online[qoff + b0_size + b1_size + a2]
|
||||
+ save_q_online[qoff + b0_size + b1_size + b2_size + a3];
|
||||
target_q[t] = save_q_target[qoff + a0]; /* target_q already contains Bellman target */
|
||||
}
|
||||
}
|
||||
@@ -623,6 +640,15 @@ void iqn_backward_kernel(
|
||||
if (tid == 0)
|
||||
atomicAdd(&grad_buf[off[7] + a2], inv_batch * dL_dq);
|
||||
|
||||
/* Branch 3 */
|
||||
for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) {
|
||||
dL_dcomb[h / IQN_BLOCK_SIZE] += w_b3[a3 * hidden_dim + h] * dL_dq;
|
||||
atomicAdd(&grad_buf[off[8] + a3 * hidden_dim + h],
|
||||
inv_batch * dL_dq * comb_dist[h / IQN_BLOCK_SIZE]);
|
||||
}
|
||||
if (tid == 0)
|
||||
atomicAdd(&grad_buf[off[9] + a3], inv_batch * dL_dq);
|
||||
|
||||
/* ── Gradient through element-wise product ── */
|
||||
/* combined = h_s2 ⊙ embed
|
||||
* dL/d(embed) = dL/d(combined) ⊙ h_s2
|
||||
@@ -671,7 +697,11 @@ void iqn_grad_norm_kernel(
|
||||
int total_params,
|
||||
int shared_h1, /* runtime: unused, for consistent interface */
|
||||
int hidden_dim, /* runtime: unused, for consistent interface */
|
||||
int embed_dim /* runtime: unused, for consistent interface */
|
||||
int embed_dim, /* runtime: unused, for consistent interface */
|
||||
int b0_size, /* runtime: unused, for consistent interface */
|
||||
int b1_size, /* runtime: unused, for consistent interface */
|
||||
int b2_size, /* runtime: unused, for consistent interface */
|
||||
int b3_size /* runtime: unused, for consistent interface */
|
||||
)
|
||||
{
|
||||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
@@ -713,7 +743,11 @@ void iqn_adam_kernel(
|
||||
int total_params,
|
||||
int shared_h1, /* runtime: unused, for consistent interface */
|
||||
int hidden_dim, /* runtime: unused, for consistent interface */
|
||||
int embed_dim /* runtime: unused, for consistent interface */
|
||||
int embed_dim, /* runtime: unused, for consistent interface */
|
||||
int b0_size, /* runtime: unused, for consistent interface */
|
||||
int b1_size, /* runtime: unused, for consistent interface */
|
||||
int b2_size, /* runtime: unused, for consistent interface */
|
||||
int b3_size /* runtime: unused, for consistent interface */
|
||||
)
|
||||
{
|
||||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
@@ -760,7 +794,7 @@ void iqn_adam_kernel(
|
||||
* Computes expected Q-values per action per branch by averaging over
|
||||
* quantile samples. Used for action selection blending with C51.
|
||||
*
|
||||
* Outputs: expected_q [B, TOTAL_BRANCH_ACTIONS] — mean Q per action.
|
||||
* Outputs: expected_q [B, tba] — mean Q per action.
|
||||
*/
|
||||
extern "C" __global__
|
||||
void iqn_forward_kernel(
|
||||
@@ -768,11 +802,15 @@ void iqn_forward_kernel(
|
||||
const float* __restrict__ taus, /* [B, N] (f32) */
|
||||
const float* __restrict__ params, /* IQN weights (f32) */
|
||||
const float* __restrict__ cos_features, /* [N, embed_dim] precomputed cosine features */
|
||||
float* __restrict__ expected_q, /* [B, TOTAL_BRANCH_ACTIONS] (f32) */
|
||||
float* __restrict__ expected_q, /* [B, tba] (f32) */
|
||||
int batch_size,
|
||||
int shared_h1, /* runtime: unused here, for consistency */
|
||||
int hidden_dim, /* runtime: IQN hidden dim = SHARED_H2 */
|
||||
int embed_dim /* runtime: cosine embedding dimension (was IQN_EMBED_DIM) */
|
||||
int embed_dim, /* runtime: cosine embedding dimension (was IQN_EMBED_DIM) */
|
||||
int b0_size, /* runtime: branch 0 (direction) action count */
|
||||
int b1_size, /* runtime: branch 1 (magnitude) action count */
|
||||
int b2_size, /* runtime: branch 2 (order) action count */
|
||||
int b3_size /* runtime: branch 3 (urgency) action count */
|
||||
)
|
||||
{
|
||||
int sample = blockIdx.x;
|
||||
@@ -783,8 +821,9 @@ void iqn_forward_kernel(
|
||||
__shared__ float shmem_reduce[IQN_BLOCK_SIZE / 32];
|
||||
|
||||
/* ── Compute runtime weight offsets ── */
|
||||
int off[8];
|
||||
iqn_compute_offsets(hidden_dim, embed_dim, off);
|
||||
int tba = b0_size + b1_size + b2_size + b3_size;
|
||||
int off[10];
|
||||
iqn_compute_offsets(hidden_dim, embed_dim, b0_size, b1_size, b2_size, b3_size, off);
|
||||
|
||||
const float* my_h_s2_bf16 = h_s2 + sample * hidden_dim;
|
||||
const float* my_taus = taus + sample * IQN_NUM_QUANTILES;
|
||||
@@ -796,15 +835,18 @@ void iqn_forward_kernel(
|
||||
const float* b_b1 = params + off[5];
|
||||
const float* w_b2 = params + off[6];
|
||||
const float* b_b2 = params + off[7];
|
||||
const float* w_b3 = params + off[8];
|
||||
const float* b_b3 = params + off[9];
|
||||
|
||||
/* Load h_s2 (bf16 → f32 at boundary) */
|
||||
float h_dist[IQN_DIST_MAX];
|
||||
for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE)
|
||||
h_dist[h / IQN_BLOCK_SIZE] = (float)my_h_s2_bf16[h];
|
||||
|
||||
/* Accumulate Q-values across quantiles (for mean) */
|
||||
float q_acc[TOTAL_BRANCH_ACTIONS];
|
||||
for (int a = 0; a < TOTAL_BRANCH_ACTIONS; a++)
|
||||
/* Accumulate Q-values across quantiles (for mean).
|
||||
* Max tba = 3+3+3+3 = 12 — fits easily in registers. */
|
||||
float q_acc[16]; /* sized >= max tba */
|
||||
for (int a = 0; a < tba; a++)
|
||||
q_acc[a] = 0.0f;
|
||||
|
||||
for (int t = 0; t < IQN_NUM_QUANTILES; t++) {
|
||||
@@ -829,7 +871,7 @@ void iqn_forward_kernel(
|
||||
|
||||
/* Branch outputs */
|
||||
/* Branch 0 */
|
||||
for (int a = 0; a < BRANCH_0_SIZE; a++) {
|
||||
for (int a = 0; a < b0_size; a++) {
|
||||
float partial = 0.0f;
|
||||
const float* w_row = w_b0 + a * hidden_dim;
|
||||
for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE)
|
||||
@@ -838,30 +880,39 @@ void iqn_forward_kernel(
|
||||
q_acc[a] += q_val;
|
||||
}
|
||||
/* Branch 1 */
|
||||
for (int a = 0; a < BRANCH_1_SIZE; a++) {
|
||||
for (int a = 0; a < b1_size; a++) {
|
||||
float partial = 0.0f;
|
||||
const float* w_row = w_b1 + a * hidden_dim;
|
||||
for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE)
|
||||
partial += w_row[h] * comb_dist[h / IQN_BLOCK_SIZE];
|
||||
float q_val = iqn_block_sum(partial, shmem_reduce) + b_b1[a];
|
||||
q_acc[BRANCH_0_SIZE + a] += q_val;
|
||||
q_acc[b0_size + a] += q_val;
|
||||
}
|
||||
/* Branch 2 */
|
||||
for (int a = 0; a < BRANCH_2_SIZE; a++) {
|
||||
for (int a = 0; a < b2_size; a++) {
|
||||
float partial = 0.0f;
|
||||
const float* w_row = w_b2 + a * hidden_dim;
|
||||
for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE)
|
||||
partial += w_row[h] * comb_dist[h / IQN_BLOCK_SIZE];
|
||||
float q_val = iqn_block_sum(partial, shmem_reduce) + b_b2[a];
|
||||
q_acc[BRANCH_0_SIZE + BRANCH_1_SIZE + a] += q_val;
|
||||
q_acc[b0_size + b1_size + a] += q_val;
|
||||
}
|
||||
/* Branch 3 */
|
||||
for (int a = 0; a < b3_size; a++) {
|
||||
float partial = 0.0f;
|
||||
const float* w_row = w_b3 + a * hidden_dim;
|
||||
for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE)
|
||||
partial += w_row[h] * comb_dist[h / IQN_BLOCK_SIZE];
|
||||
float q_val = iqn_block_sum(partial, shmem_reduce) + b_b3[a];
|
||||
q_acc[b0_size + b1_size + b2_size + a] += q_val;
|
||||
}
|
||||
}
|
||||
|
||||
/* Write mean Q-values (averaged over quantiles) */
|
||||
float inv_n = 1.0f / (float)IQN_NUM_QUANTILES;
|
||||
if (tid == 0) {
|
||||
for (int a = 0; a < TOTAL_BRANCH_ACTIONS; a++)
|
||||
expected_q[sample * TOTAL_BRANCH_ACTIONS + a] = q_acc[a] * inv_n;
|
||||
for (int a = 0; a < tba; a++)
|
||||
expected_q[sample * tba + a] = q_acc[a] * inv_n;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -891,7 +942,11 @@ void iqn_trunk_forward_kernel(
|
||||
int state_dim, /* runtime state dimension (replaces IQN_STATE_DIM) */
|
||||
int shared_h1, /* runtime: first hidden layer width */
|
||||
int hidden_dim, /* runtime: IQN hidden dim = SHARED_H2 */
|
||||
int embed_dim /* runtime: unused here, for consistent interface */
|
||||
int embed_dim, /* runtime: unused here, for consistent interface */
|
||||
int b0_size, /* runtime: unused here, for consistent interface */
|
||||
int b1_size, /* runtime: unused here, for consistent interface */
|
||||
int b2_size, /* runtime: unused here, for consistent interface */
|
||||
int b3_size /* runtime: unused here, for consistent interface */
|
||||
)
|
||||
{
|
||||
int sample = blockIdx.x;
|
||||
@@ -941,7 +996,11 @@ void iqn_ema_kernel(
|
||||
int n,
|
||||
int shared_h1,
|
||||
int hidden_dim,
|
||||
int embed_dim
|
||||
int embed_dim,
|
||||
int b0_size,
|
||||
int b1_size,
|
||||
int b2_size,
|
||||
int b3_size
|
||||
)
|
||||
{
|
||||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
@@ -953,11 +1012,12 @@ void iqn_ema_kernel(
|
||||
|
||||
/* ── Flat action → branch action decode kernel ─────────────────────────
|
||||
*
|
||||
* Decodes flat action index ∈ [0, B0×B1×B2) into 3 branch indices.
|
||||
* Decodes flat action index ∈ [0, b0×b1×b2×b3) into 4 branch indices.
|
||||
* action = dir*b1*b2*b3 + mag*b2*b3 + order*b3 + urgency
|
||||
* Grid: (ceil(batch_size/256), 1, 1), Block: (256, 1, 1)
|
||||
*
|
||||
* Input: flat_actions[B] -- i32 flat action indices
|
||||
* Output: branch_actions[B*3] -- i32 [exposure, order, urgency] per sample
|
||||
* Output: branch_actions[B*4] -- i32 [direction, magnitude, order, urgency] per sample
|
||||
*/
|
||||
extern "C" __global__
|
||||
void iqn_decode_actions_kernel(
|
||||
@@ -966,7 +1026,11 @@ void iqn_decode_actions_kernel(
|
||||
int batch_size,
|
||||
int shared_h1, /* runtime: unused, for consistent interface */
|
||||
int hidden_dim, /* runtime: unused, for consistent interface */
|
||||
int embed_dim /* runtime: unused, for consistent interface */
|
||||
int embed_dim, /* runtime: unused, for consistent interface */
|
||||
int b0_size, /* runtime: branch 0 (direction) action count */
|
||||
int b1_size, /* runtime: branch 1 (magnitude) action count */
|
||||
int b2_size, /* runtime: branch 2 (order) action count */
|
||||
int b3_size /* runtime: branch 3 (urgency) action count */
|
||||
)
|
||||
{
|
||||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
@@ -974,17 +1038,21 @@ void iqn_decode_actions_kernel(
|
||||
|
||||
int flat = flat_actions[i];
|
||||
/* Bounds check: clamp to valid action range [0, total_actions) */
|
||||
int max_action = BRANCH_0_SIZE * BRANCH_1_SIZE * BRANCH_2_SIZE;
|
||||
int max_action = b0_size * b1_size * b2_size * b3_size;
|
||||
if (flat < 0 || flat >= max_action) flat = 0;
|
||||
int stride = BRANCH_1_SIZE * BRANCH_2_SIZE;
|
||||
int exposure = flat / stride;
|
||||
int remainder = flat % stride;
|
||||
int order = remainder / BRANCH_2_SIZE;
|
||||
int urgency = remainder % BRANCH_2_SIZE;
|
||||
int s1 = b1_size * b2_size * b3_size;
|
||||
int s2 = b2_size * b3_size;
|
||||
int direction = flat / s1;
|
||||
int remainder = flat % s1;
|
||||
int magnitude = remainder / s2;
|
||||
remainder = remainder % s2;
|
||||
int order = remainder / b3_size;
|
||||
int urgency = remainder % b3_size;
|
||||
|
||||
branch_actions[i * 3 + 0] = exposure;
|
||||
branch_actions[i * 3 + 1] = order;
|
||||
branch_actions[i * 3 + 2] = urgency;
|
||||
branch_actions[i * 4 + 0] = direction;
|
||||
branch_actions[i * 4 + 1] = magnitude;
|
||||
branch_actions[i * 4 + 2] = order;
|
||||
branch_actions[i * 4 + 3] = urgency;
|
||||
}
|
||||
|
||||
/* ── GPU τ sampling kernel (Philox-based PRNG) ─────────────────────────
|
||||
@@ -1022,7 +1090,11 @@ void iqn_sample_taus_kernel(
|
||||
int total,
|
||||
int shared_h1, /* runtime: unused, for consistent interface */
|
||||
int hidden_dim, /* runtime: unused, for consistent interface */
|
||||
int embed_dim /* runtime: unused, for consistent interface */
|
||||
int embed_dim, /* runtime: unused, for consistent interface */
|
||||
int b0_size, /* runtime: unused, for consistent interface */
|
||||
int b1_size, /* runtime: unused, for consistent interface */
|
||||
int b2_size, /* runtime: unused, for consistent interface */
|
||||
int b3_size /* runtime: unused, for consistent interface */
|
||||
)
|
||||
{
|
||||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
@@ -50,11 +50,12 @@ extern "C" __global__ void mse_grad_kernel(
|
||||
float inv_batch = 1.0f / (float)batch_size;
|
||||
float d_combined = inv_batch * isw * td_error * p_j * (z_j - e_q);
|
||||
|
||||
/* Magnitude entropy boost: prevent atom distribution collapse for d==1.
|
||||
* MSE path has no entropy param — hardcoded 0.005 = 5× base (0.001). */
|
||||
if (d == 1) {
|
||||
/* Per-branch entropy boost: prevent distribution collapse.
|
||||
* Magnitude (d==1): 0.005 = 5× base. Order (d==2): 0.003 = 3× base. */
|
||||
if (d == 1 || d == 2) {
|
||||
float ent_weight = (d == 1) ? 0.005f : 0.003f;
|
||||
float lp_approx = fmaxf(logf(fmaxf(p_j, 1e-8f)), -10.0f);
|
||||
d_combined += 0.005f * (1.0f + lp_approx);
|
||||
d_combined += ent_weight * (1.0f + lp_approx);
|
||||
}
|
||||
|
||||
/* Route through dueling: d_value[b,j] += d_combined.
|
||||
|
||||
548
docs/superpowers/plans/2026-04-11-action-diversity-iqn4branch.md
Normal file
548
docs/superpowers/plans/2026-04-11-action-diversity-iqn4branch.md
Normal file
@@ -0,0 +1,548 @@
|
||||
# Action Diversity + IQN 4-Branch Fix Implementation Plan
|
||||
|
||||
> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
|
||||
|
||||
**Goal:** Fix IQN 3→4 branch kernel mismatch (corrupting 75% of gradient budget), add order-type counterfactual relabeling, per-branch entropy regularization, and position histogram for all 4 branches — raising order diversity from 1/3 to 3/3.
|
||||
|
||||
**Architecture:** Five compounding causes of order branch collapse, fixed bottom-up: (1) IQN kernel uses pre-4-branch 3-branch defines with BRANCH_0_SIZE=5, corrupting 2/3 of gradient signal; (2) no counterfactual experience for order types; (3) no per-branch entropy boost for order; (4) position histogram only tracks dir+mag; (5) Flat-state trains order Q-values on null transitions. All changes are in CUDA kernels + build.rs + Rust config — no CUDA graph structure changes.
|
||||
|
||||
**Tech Stack:** CUDA 12.4 (nvcc cubin), Rust 1.85, cudarc
|
||||
|
||||
---
|
||||
|
||||
## File Map
|
||||
|
||||
| File | Action | Responsibility |
|
||||
|------|--------|---------------|
|
||||
| `crates/ml/build.rs` | Modify | Add `-D` defines for IQN branch sizes |
|
||||
| `crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu` | Modify | 3→4 branch: defines, offsets, forward, backward, decode |
|
||||
| `crates/ml/src/cuda_pipeline/gpu_iqn_head.rs` | Modify | Update buffer sizes for 4 branches, fix `total_branch_actions` in save buffers |
|
||||
| `crates/ml/src/cuda_pipeline/experience_kernels.cu` | Modify | 3-cycle counterfactual + order/urgency histogram |
|
||||
| `crates/ml/src/cuda_pipeline/c51_grad_kernel.cu` | Modify | Per-branch entropy boost for order (d==2) |
|
||||
| `crates/ml/src/cuda_pipeline/mse_grad_kernel.cu` | Modify | Per-branch entropy boost for order (d==2) |
|
||||
|
||||
---
|
||||
|
||||
### Task 1: IQN kernel — runtime branch sizes, eliminate compile-time defines
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu:60-99,188-199`
|
||||
|
||||
No build.rs changes needed — branch sizes become runtime kernel parameters like `hidden_dim` and `embed_dim` already are. This prevents the entire class of stale-define bugs (the compile-time BRANCH_0_SIZE=5 was never updated when 4-branch landed).
|
||||
|
||||
- [ ] **Step 1: Remove compile-time BRANCH_*_SIZE defines, add runtime parameters**
|
||||
|
||||
In `iqn_dual_head_kernel.cu`, delete the compile-time branch defines (lines 62-72) and the `IQN_W_B*_SIZE` / `IQN_TOTAL_PARAMS` macros (lines 83-95). Replace with a comment:
|
||||
|
||||
```c
|
||||
/* Branch sizes are RUNTIME parameters — passed to every kernel.
|
||||
* Eliminates stale compile-time define bugs (BRANCH_0_SIZE was 5 when 4-branch uses 3).
|
||||
* Weight layout: W_embed[H,D], b_embed[H], then per-branch W_bd[Bd,H], b_bd[Bd]. */
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Update offset function to use runtime branch sizes**
|
||||
|
||||
Replace `iqn_compute_offsets` (lines 188-199) to accept runtime branch sizes:
|
||||
|
||||
```c
|
||||
/** Compute runtime weight offsets from hidden_dim, embed_dim, and 4 branch sizes.
|
||||
* Returns offsets[10]: W_EMBED, B_EMBED, W_B0, B_B0, W_B1, B_B1, W_B2, B_B2, W_B3, B_B3 */
|
||||
__device__ __forceinline__ void iqn_compute_offsets(
|
||||
int hidden_dim, int embed_dim,
|
||||
int b0, int b1, int b2, int b3,
|
||||
int offsets[10]
|
||||
) {
|
||||
offsets[0] = 0; /* W_EMBED */
|
||||
offsets[1] = offsets[0] + hidden_dim * embed_dim; /* B_EMBED */
|
||||
offsets[2] = offsets[1] + hidden_dim; /* W_B0 */
|
||||
offsets[3] = offsets[2] + b0 * hidden_dim; /* B_B0 */
|
||||
offsets[4] = offsets[3] + b0; /* W_B1 */
|
||||
offsets[5] = offsets[4] + b1 * hidden_dim; /* B_B1 */
|
||||
offsets[6] = offsets[5] + b1; /* W_B2 */
|
||||
offsets[7] = offsets[6] + b2 * hidden_dim; /* B_B2 */
|
||||
offsets[8] = offsets[7] + b2; /* W_B3 */
|
||||
offsets[9] = offsets[8] + b3 * hidden_dim; /* B_B3 */
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Add branch size parameters to all 9 kernel signatures**
|
||||
|
||||
Every IQN kernel signature gets 4 new `int` parameters: `int b0_size, int b1_size, int b2_size, int b3_size`. Replace all internal uses of `BRANCH_0_SIZE` etc. with these runtime params. Also compute `int tba = b0_size + b1_size + b2_size + b3_size;` at the top of each kernel (replaces `TOTAL_BRANCH_ACTIONS`).
|
||||
|
||||
For `iqn_forward_loss_kernel`, `iqn_backward_kernel`, `iqn_forward_kernel`:
|
||||
```c
|
||||
/* Add after existing params (hidden_dim, embed_dim): */
|
||||
int b0_size, int b1_size, int b2_size, int b3_size
|
||||
```
|
||||
|
||||
For `iqn_decode_actions_kernel`, `iqn_grad_norm_kernel`, `iqn_adam_kernel`, `iqn_ema_kernel`, `iqn_trunk_forward_kernel`, `iqn_sample_taus_kernel`: add the same 4 params even if unused (consistent interface, like `shared_h1` is already passed to some kernels unused).
|
||||
|
||||
- [ ] **Step 4: Update Rust kernel launch calls in `gpu_iqn_head.rs`**
|
||||
|
||||
Every `.launch_builder()` call for IQN kernels needs 4 additional `.arg()` calls:
|
||||
```rust
|
||||
.arg(&(self.config.branch_0_size as i32))
|
||||
.arg(&(self.config.branch_1_size as i32))
|
||||
.arg(&(self.config.branch_2_size as i32))
|
||||
.arg(&(self.config.branch_3_size as i32))
|
||||
```
|
||||
|
||||
- [ ] **Step 5: Verify kernel compiles and all tests pass**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo build -p ml --lib 2>&1 | grep -E "iqn_dual_head|error|FAILED"
|
||||
SQLX_OFFLINE=true cargo check -p ml --lib
|
||||
```
|
||||
Expected: `Compiled iqn_dual_head_kernel.cu -> iqn_dual_head_kernel.cubin` with no errors.
|
||||
|
||||
- [ ] **Step 6: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu crates/ml/src/cuda_pipeline/gpu_iqn_head.rs
|
||||
git commit -m "fix(iqn): runtime branch sizes — eliminate stale compile-time defines"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 2: IQN kernel — 4-branch forward, backward, and decode
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu:215-475,488-664,766-988`
|
||||
|
||||
- [ ] **Step 1: Update decode kernel to 4-branch**
|
||||
|
||||
Replace `iqn_decode_actions_kernel` (lines 962-988). Output becomes `[B, 4]` (was `[B, 3]`):
|
||||
|
||||
```c
|
||||
/* Output: branch_actions[B*4] -- i32 [dir, order, urgency, magnitude] per sample */
|
||||
extern "C" __global__
|
||||
void iqn_decode_actions_kernel(
|
||||
const int* __restrict__ flat_actions,
|
||||
int* __restrict__ branch_actions,
|
||||
int batch_size,
|
||||
int shared_h1,
|
||||
int hidden_dim,
|
||||
int embed_dim
|
||||
)
|
||||
{
|
||||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (i >= batch_size) return;
|
||||
|
||||
int flat = flat_actions[i];
|
||||
int max_action = BRANCH_0_SIZE * BRANCH_1_SIZE * BRANCH_2_SIZE * BRANCH_3_SIZE;
|
||||
if (flat < 0 || flat >= max_action) flat = 0;
|
||||
|
||||
/* 4-branch decode: action = dir*b1*b2*b3 + mag*b2*b3 + order*b3 + urgency */
|
||||
int stride_0 = BRANCH_1_SIZE * BRANCH_2_SIZE * BRANCH_3_SIZE;
|
||||
int stride_1 = BRANCH_2_SIZE * BRANCH_3_SIZE;
|
||||
int stride_2 = BRANCH_3_SIZE;
|
||||
|
||||
int dir = flat / stride_0;
|
||||
int rem = flat % stride_0;
|
||||
int mag = rem / stride_1;
|
||||
rem = rem % stride_1;
|
||||
int order = rem / stride_2;
|
||||
int urgency = rem % stride_2;
|
||||
|
||||
branch_actions[i * 4 + 0] = dir;
|
||||
branch_actions[i * 4 + 1] = mag;
|
||||
branch_actions[i * 4 + 2] = order;
|
||||
branch_actions[i * 4 + 3] = urgency;
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Update forward+loss kernel for 4 branches**
|
||||
|
||||
In `iqn_forward_loss_kernel`, add branch 3 to the online forward (after branch 2 block at line ~380), target forward (after line ~437), and sum of taken Q-values. The pattern is identical to branches 0-2 — just add:
|
||||
|
||||
```c
|
||||
/* Branch 3 (magnitude) — online */
|
||||
float q3_taken = 0.0f;
|
||||
for (int a = 0; a < BRANCH_3_SIZE; a++) {
|
||||
float partial = 0.0f;
|
||||
const float* w_row = w_b3 + a * hidden_dim;
|
||||
for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE)
|
||||
partial += w_row[h] * comb_dist[h / IQN_BLOCK_SIZE];
|
||||
float acc = b_b3[a] + iqn_block_sum(partial, shmem_reduce);
|
||||
if (tid == 0) {
|
||||
save_q_online[q_save_offset + BRANCH_0_SIZE + BRANCH_1_SIZE + BRANCH_2_SIZE + a] = acc;
|
||||
if (a == a3) q3_taken = acc;
|
||||
}
|
||||
}
|
||||
if (tid == 0) shmem_reduce[0] = q3_taken;
|
||||
__syncthreads();
|
||||
q3_taken = shmem_reduce[0];
|
||||
```
|
||||
|
||||
Update the sum line: `online_q_a[t] = q0_taken + q1_taken + q2_taken + q3_taken;`
|
||||
|
||||
Similarly for the target branch and `target_q_a` save.
|
||||
|
||||
Also update action decode at the top of the kernel:
|
||||
```c
|
||||
int a0 = actions[sample * 4 + 0]; // was sample * 3
|
||||
int a1 = actions[sample * 4 + 1];
|
||||
int a2 = actions[sample * 4 + 2];
|
||||
int a3 = actions[sample * 4 + 3];
|
||||
```
|
||||
|
||||
And add weight pointers:
|
||||
```c
|
||||
const float* w_b3 = online_params + off[8];
|
||||
const float* b_b3 = online_params + off[9];
|
||||
const float* tw_b3 = target_params + off[8];
|
||||
const float* tb_b3 = target_params + off[9];
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Update backward kernel for 4 branches**
|
||||
|
||||
In `iqn_backward_kernel`, add branch 3 gradient computation after the branch 2 block (after line ~624). Pattern identical:
|
||||
|
||||
```c
|
||||
/* Branch 3 (magnitude) */
|
||||
for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) {
|
||||
dL_dcomb[h / IQN_BLOCK_SIZE] += w_b3[a3 * hidden_dim + h] * dL_dq;
|
||||
atomicAdd(&grad_buf[off[8] + a3 * hidden_dim + h],
|
||||
inv_batch * dL_dq * comb_dist[h / IQN_BLOCK_SIZE]);
|
||||
}
|
||||
if (tid == 0)
|
||||
atomicAdd(&grad_buf[off[9] + a3], inv_batch * dL_dq);
|
||||
```
|
||||
|
||||
Update action reads: `int a3 = actions[sample * 4 + 3];` and add `const float* w_b3 = online_params + off[8];`
|
||||
|
||||
Update offset array: `int off[10]; iqn_compute_offsets(hidden_dim, embed_dim, off);`
|
||||
|
||||
- [ ] **Step 4: Update forward-only (inference) kernel for 4 branches**
|
||||
|
||||
The `iqn_forward_kernel` (line 766) also needs branch 3. Same pattern as forward+loss but without loss computation. Add branch 3 Q-value computation and update `q_total += q3`.
|
||||
|
||||
- [ ] **Step 5: Verify kernel compiles**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo build -p ml --lib 2>&1 | grep -E "iqn_dual_head|error"
|
||||
```
|
||||
|
||||
- [ ] **Step 6: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu
|
||||
git commit -m "fix(iqn): 4-branch forward, backward, decode — was 3-branch corrupting 75% gradient"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 3: IQN Rust side — update buffer sizes and decode for 4 branches
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_iqn_head.rs:113-127,233-293,339-345`
|
||||
|
||||
- [ ] **Step 1: Update `branch_actions` allocation from `B*3` to `B*4`**
|
||||
|
||||
In `GpuIqnHead::new()`, change:
|
||||
```rust
|
||||
// Was: alloc_zeros::<i32>(b * 3)
|
||||
let branch_actions = stream.alloc_zeros::<i32>(b * 4).map_err(|e| {
|
||||
MLError::ModelError(format!("IQN alloc branch_actions: {e}"))
|
||||
})?;
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Update `save_q_online` and `save_q_target` buffer sizes**
|
||||
|
||||
These use `tba = config.total_branch_actions()` which already returns `b0+b1+b2+b3=12`. Verify `tba` is used (not hardcoded). Check lines 292-293:
|
||||
```rust
|
||||
let save_q_online = alloc_f32(&stream, b * n * tba, "iqn_save_q_online")?;
|
||||
let save_q_target = alloc_f32(&stream, b * n * tba, "iqn_save_q_target")?;
|
||||
```
|
||||
These should already be correct since `total_branch_actions()` includes b3.
|
||||
|
||||
- [ ] **Step 3: Update IQR buffer size**
|
||||
|
||||
Line 306: `let iqr_buf = alloc_f32(&stream, tba, "iqn_iqr")?;` — already uses `tba=12`. OK.
|
||||
|
||||
- [ ] **Step 4: Verify all IQN-related tests pass**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo check -p ml --lib
|
||||
SQLX_OFFLINE=true cargo test -p ml --lib -- gradient_budget --nocapture
|
||||
```
|
||||
|
||||
- [ ] **Step 5: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/cuda_pipeline/gpu_iqn_head.rs
|
||||
git commit -m "fix(iqn): update Rust buffer sizes for 4-branch decode (B*4 actions)"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 4: Order-type counterfactual relabeling (3-cycle)
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/experience_kernels.cu:1799-1846`
|
||||
|
||||
- [ ] **Step 1: Replace 2-cycle counterfactual with 3-cycle**
|
||||
|
||||
Replace the `if ((current_t & 1) == 0)` block (lines 1799-1846) with a 3-cycle that adds order-type counterfactuals:
|
||||
|
||||
```c
|
||||
int cf_cycle = current_t % 3;
|
||||
|
||||
if (cf_cycle == 0) {
|
||||
/* ── Directional mirror (existing) ── */
|
||||
int cf_dir = (b0_size - 1) - dir_idx;
|
||||
cf_action = cf_dir * b1_size * b2_size * b3_size
|
||||
+ mag_idx * b2_size * b3_size
|
||||
+ orig_order * b3_size + orig_urgency;
|
||||
cf_reward = -reward;
|
||||
} else if (cf_cycle == 1) {
|
||||
/* ── Hindsight magnitude relabeling (existing) ── */
|
||||
const float mag_fracs[3] = { 0.50f, 0.707f, 1.00f };
|
||||
float actual_mag_frac = mag_fracs[mag_idx < 3 ? mag_idx : 0];
|
||||
|
||||
if (fabsf(reward) > 0.001f) {
|
||||
int alt_mag;
|
||||
int cycle = (current_t / 3) % 3;
|
||||
if (cycle == 0) alt_mag = (mag_idx > 0) ? mag_idx - 1 : 2;
|
||||
else if (cycle == 1) alt_mag = (mag_idx < 2) ? mag_idx + 1 : 0;
|
||||
else alt_mag = 2 - mag_idx;
|
||||
|
||||
float alt_mag_frac = mag_fracs[alt_mag];
|
||||
cf_action = dir_idx * b1_size * b2_size * b3_size
|
||||
+ alt_mag * b2_size * b3_size
|
||||
+ orig_order * b3_size + orig_urgency;
|
||||
cf_reward = reward * (alt_mag_frac / actual_mag_frac);
|
||||
cf_reward = fminf(fmaxf(cf_reward, -10.0f), 10.0f);
|
||||
} else {
|
||||
int cf_dir = (b0_size - 1) - dir_idx;
|
||||
cf_action = cf_dir * b1_size * b2_size * b3_size
|
||||
+ mag_idx * b2_size * b3_size
|
||||
+ orig_order * b3_size + orig_urgency;
|
||||
cf_reward = -reward;
|
||||
}
|
||||
} else {
|
||||
/* ── NEW: Order-type counterfactual ──
|
||||
* Cycle through alternative order types. Compute execution cost
|
||||
* difference and amplify it to be comparable to directional signal.
|
||||
* Only meaningful when position is non-zero (order type affects cost). */
|
||||
int alt_order = (orig_order + 1 + ((current_t / 3) & 1)) % b2_size;
|
||||
|
||||
if (fabsf(position) > 0.001f) {
|
||||
/* Compute cost delta between taken and alternative order type */
|
||||
float taken_cost = compute_tx_cost(
|
||||
position - pre_trade_position, raw_close, tx_cost_multiplier,
|
||||
spread_cost, max_position, orig_order, spread_scale);
|
||||
float alt_cost = compute_tx_cost(
|
||||
position - pre_trade_position, raw_close, tx_cost_multiplier,
|
||||
spread_cost, max_position, alt_order, spread_scale);
|
||||
float cost_delta = (taken_cost - alt_cost)
|
||||
/ fmaxf(prev_equity, 1.0f);
|
||||
|
||||
cf_action = dir_idx * b1_size * b2_size * b3_size
|
||||
+ mag_idx * b2_size * b3_size
|
||||
+ alt_order * b3_size + orig_urgency;
|
||||
/* Amplify 50×: raw cost delta ~0.0001-0.0005, directional ~0.01-0.1 */
|
||||
cf_reward = reward + 50.0f * cost_delta;
|
||||
cf_reward = fminf(fmaxf(cf_reward, -10.0f), 10.0f);
|
||||
} else {
|
||||
/* Flat position: order type irrelevant, fall back to directional mirror */
|
||||
int cf_dir = (b0_size - 1) - dir_idx;
|
||||
cf_action = cf_dir * b1_size * b2_size * b3_size
|
||||
+ mag_idx * b2_size * b3_size
|
||||
+ orig_order * b3_size + orig_urgency;
|
||||
cf_reward = -reward;
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Note: `position`, `pre_trade_position`, `raw_close`, `tx_cost_multiplier`, `spread_cost`, `max_position`, `spread_scale` are all already in scope from the experience_env_step kernel's local variables. `compute_tx_cost` is declared in `trade_physics.cuh` which is already included.
|
||||
|
||||
- [ ] **Step 2: Verify kernel compiles**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo build -p ml --lib 2>&1 | grep -E "experience_kernels|error"
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/cuda_pipeline/experience_kernels.cu
|
||||
git commit -m "feat: order-type counterfactual relabeling (3-cycle: dir/mag/order)"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 5: Per-branch entropy regularization for order
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/c51_grad_kernel.cu:53-57`
|
||||
- Modify: `crates/ml/src/cuda_pipeline/mse_grad_kernel.cu:53-58`
|
||||
|
||||
- [ ] **Step 1: Add order entropy boost to C51 grad kernel**
|
||||
|
||||
Replace lines 53-57 in `c51_grad_kernel.cu`:
|
||||
|
||||
```c
|
||||
/* Entropy regularization — per-branch coefficients.
|
||||
* Magnitude (d==1): 5× boost for atom distribution collapse prevention.
|
||||
* Order (d==2): 3× boost to prevent order type collapse.
|
||||
* Direction (d==0), urgency (d==3): baseline 1×. */
|
||||
if (entropy_coeff > 0.0f) {
|
||||
float ent_scale;
|
||||
if (d == 1) ent_scale = 5.0f;
|
||||
else if (d == 2) ent_scale = 3.0f;
|
||||
else ent_scale = 1.0f;
|
||||
float lp_clamped = fmaxf(lp, -10.0f);
|
||||
d_combined += ent_scale * entropy_coeff * (1.0f + lp_clamped);
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Add order entropy boost to MSE grad kernel**
|
||||
|
||||
Replace lines 53-58 in `mse_grad_kernel.cu`:
|
||||
|
||||
```c
|
||||
/* Per-branch entropy boost: prevent distribution collapse.
|
||||
* Magnitude (d==1): 0.005 = 5× base (0.001). Order (d==2): 0.003 = 3× base. */
|
||||
if (d == 1 || d == 2) {
|
||||
float ent_weight = (d == 1) ? 0.005f : 0.003f;
|
||||
float lp_approx = fmaxf(logf(fmaxf(p_j, 1e-8f)), -10.0f);
|
||||
d_combined += ent_weight * (1.0f + lp_approx);
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Verify kernel compiles**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo build -p ml --lib 2>&1 | grep -E "c51_grad\|mse_grad|error"
|
||||
```
|
||||
|
||||
- [ ] **Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/cuda_pipeline/c51_grad_kernel.cu crates/ml/src/cuda_pipeline/mse_grad_kernel.cu
|
||||
git commit -m "feat: per-branch entropy regularization — order gets 3× boost"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 6: Position histogram for all 4 branches
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/experience_kernels.cu:1664-1707`
|
||||
|
||||
- [ ] **Step 1: Expand histogram from 6 bins to 12 bins (3×4 branches)**
|
||||
|
||||
Replace the histogram tracking block (lines 1664-1707):
|
||||
|
||||
```c
|
||||
/* ── #19 Position entropy: track all 4 branches (3 bins each = 12 total) ── */
|
||||
if (position_histogram != NULL && position_entropy_weight > 0.0f) {
|
||||
/* Histogram layout: [N, 12] — 3 bins per branch: dir[0:3], mag[3:6], order[6:9], urgency[9:12] */
|
||||
int order_idx = decode_order_4b(action_idx, b2_size, b3_size);
|
||||
int urgency_idx = decode_urgency_4b(action_idx, b3_size);
|
||||
if (dir_idx >= 0 && dir_idx < 3)
|
||||
position_histogram[(long long)i * 12 + dir_idx] += 1.0f;
|
||||
if (mag_idx >= 0 && mag_idx < 3)
|
||||
position_histogram[(long long)i * 12 + 3 + mag_idx] += 1.0f;
|
||||
if (order_idx >= 0 && order_idx < 3)
|
||||
position_histogram[(long long)i * 12 + 6 + order_idx] += 1.0f;
|
||||
if (urgency_idx >= 0 && urgency_idx < 3)
|
||||
position_histogram[(long long)i * 12 + 9 + urgency_idx] += 1.0f;
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Update entropy bonus computation to cover all 4 branches**
|
||||
|
||||
Replace the done-block entropy bonus (lines 1679-1707):
|
||||
|
||||
```c
|
||||
/* #19 Position entropy bonus at episode end — all 4 branches */
|
||||
if (done && position_histogram != NULL && position_entropy_weight > 0.0f) {
|
||||
float* hist = position_histogram + (long long)i * 12;
|
||||
float combined_entropy = 0.0f;
|
||||
/* Compute entropy over each branch's 3 bins */
|
||||
for (int branch = 0; branch < 4; branch++) {
|
||||
float* branch_hist = hist + branch * 3;
|
||||
float total = 0.0f;
|
||||
for (int b = 0; b < 3; b++) total += branch_hist[b];
|
||||
float ent = 0.0f;
|
||||
if (total > 1.0f) {
|
||||
for (int b = 0; b < 3; b++) {
|
||||
float p = branch_hist[b] / total;
|
||||
if (p > 1e-6f) ent -= p * logf(p);
|
||||
}
|
||||
ent /= 1.0986f; /* Normalize by log(3) → [0, 1] */
|
||||
}
|
||||
combined_entropy += ent;
|
||||
}
|
||||
/* Average across all 4 branches (was 2) */
|
||||
combined_entropy *= 0.25f;
|
||||
reward += position_entropy_weight * combined_entropy;
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Update position_histogram allocation in Rust**
|
||||
|
||||
Search for `position_histogram` allocation in `gpu_experience_collector.rs` and change from `alloc_episodes * 6` to `alloc_episodes * 12`:
|
||||
|
||||
```rust
|
||||
// Was: alloc_episodes * 6
|
||||
let position_histogram = stream.alloc_zeros::<f32>(alloc_episodes * 12)
|
||||
.map_err(|e| MLError::ModelError(format!("alloc position_histogram: {e}")))?;
|
||||
```
|
||||
|
||||
- [ ] **Step 4: Verify compilation and run tests**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo check -p ml --lib
|
||||
SQLX_OFFLINE=true cargo test -p ml-dqn --lib
|
||||
SQLX_OFFLINE=true cargo test -p ml --lib -- gradient_budget config monitoring
|
||||
```
|
||||
|
||||
- [ ] **Step 5: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/cuda_pipeline/experience_kernels.cu crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
|
||||
git commit -m "feat: position histogram for all 4 branches — order+urgency get entropy bonus"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 7: Full integration test — smoke test + compute-sanitizer
|
||||
|
||||
**Files:** (no changes — verification only)
|
||||
|
||||
- [ ] **Step 1: Run full smoke test**
|
||||
|
||||
```bash
|
||||
FOXHUNT_TEST_DATA=test_data/futures-baseline SQLX_OFFLINE=true cargo test -p ml --lib -- smoke_tests::feature_coverage --ignored --test-threads=1 --nocapture 2>&1 | grep -E "Action diversity|order|entropy|LOW ORDER"
|
||||
```
|
||||
|
||||
Expected: `Action diversity` should show order > 1/3. `LOW ORDER DIVERSITY` warnings should disappear or reduce.
|
||||
|
||||
- [ ] **Step 2: Run compute-sanitizer**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo test -p ml --lib --no-run 2>&1 | tail -1
|
||||
FOXHUNT_TEST_DATA=test_data/futures-baseline compute-sanitizer --tool memcheck target/debug/deps/ml-*.test "smoke_tests::feature_coverage" --ignored --test-threads=1
|
||||
```
|
||||
|
||||
Expected: `ERROR SUMMARY: 0 errors`
|
||||
|
||||
- [ ] **Step 3: Run full test suite**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo test -p ml-dqn --lib
|
||||
SQLX_OFFLINE=true cargo test -p ml --lib
|
||||
FOXHUNT_TEST_DATA=test_data/futures-baseline SQLX_OFFLINE=true cargo test -p ml --lib -- smoke_tests --ignored --test-threads=1
|
||||
```
|
||||
|
||||
Expected: All pass (286 + 903 + 19).
|
||||
|
||||
- [ ] **Step 4: Final commit with all changes**
|
||||
|
||||
```bash
|
||||
git add -A
|
||||
git commit -m "fix: IQN 4-branch + order counterfactual + per-branch entropy — action diversity restored"
|
||||
```
|
||||
Reference in New Issue
Block a user