From ca6e0007dad85240f15fd826ae6d2e16b7bc37b2 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 2 May 2026 01:36:48 +0200 Subject: [PATCH] =?UTF-8?q?docs(sp6):=20brainstorm=20spec=20=E2=80=94=20pe?= =?UTF-8?q?r-branch=20consumer=20infrastructure=20refactor?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Sonnet 4.6 --- ...branch-consumer-infrastructure-refactor.md | 297 ++++++++++++++++++ 1 file changed, 297 insertions(+) create mode 100644 docs/superpowers/specs/2026-05-02-sp6-per-branch-consumer-infrastructure-refactor.md diff --git a/docs/superpowers/specs/2026-05-02-sp6-per-branch-consumer-infrastructure-refactor.md b/docs/superpowers/specs/2026-05-02-sp6-per-branch-consumer-infrastructure-refactor.md new file mode 100644 index 000000000..2ec4934f2 --- /dev/null +++ b/docs/superpowers/specs/2026-05-02-sp6-per-branch-consumer-infrastructure-refactor.md @@ -0,0 +1,297 @@ +# SP6: Per-Branch Consumer Infrastructure Refactor + +**Date:** 2026-05-02 +**Status:** Brainstorm — open questions listed; not a plan +**Branch:** sp5-magnitude-differentiation (brainstorm only; production branch TBD) +**Upstream:** SP5 HEAD (commit `3ad5e011b` Layer B fix-up) + +--- + +## 1. Background and motivation + +SP5 Layer B (`99367b9c6`) migrated 11 consumers from hardcoded multipliers to ISV-driven adaptive signals. Eight of those consumers migrated cleanly at full per-branch fidelity. Three could not: the existing consumer infrastructure only accepts a single scalar, so each of the three consumers collapses its 4-slot per-branch ISV read to a scalar mean before passing it downstream. + +The three structural deviations are documented at the SP5 Layer B wire-up audit commit (`3ad5e011b`): + +- **Pearl 2 (loss budgets)** — `compute_adaptive_budgets()` at `crates/ml/src/trainers/dqn/fused_training.rs:3295` averages `ISV[BUDGET_C51_BASE..+4]`, `ISV[BUDGET_IQN_BASE..+4]`, `ISV[BUDGET_CQL_BASE..+4]`, `ISV[BUDGET_ENS_BASE..+4]` into four scalar scalars, then passes each to a single-broadcast `apply_c51_budget_scale(c51_budget: f32)` / `apply_cql_saxpy(cql_budget: f32)` / `apply_iqn_trunk_gradient(..., iqn_budget: f32)`. The SAXPY operations add to the full flat `grad_buf` spanning all 163 weight tensors — there is no per-branch dispatch. + +- **Pearl 3 (NoisyNet σ)** — `training_loop.rs:1747` averages `ISV[NOISY_SIGMA_BASE..+4]` into a single `mean.max(0.01)`, which is passed as `ExperienceCollectorConfig.noise_sigma: f32` (defined at `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs:298`). One scalar flows into `experience_kernels.cu` (called at `gpu_experience_collector.rs:3640`), applying the same Gaussian noise magnitude to all branches. + +- **Pearl 5 (IQN τ schedule)** — `refresh_taus_from_isv()` at `crates/ml/src/cuda_pipeline/gpu_iqn_head.rs:1962` averages the 4-branch × 5-quantile `ISV[250..270)` block per quantile slot, producing a single `[5]` schedule tiled into `[B, 5]` tensors for all IQN heads. + +The cost of these deviations is that the signal that SP5 invested ~5500 LOC to produce per-branch is discarded at the point of consumption. The producer side is complete and correct; the consumer side is not yet capable of receiving it. + +### Why it matters: the original findings + +From the 50-epoch baseline (commit `a60e7b092`) that motivated SP5: + +| Finding | Producer fix (SP5) | Consumer gap (SP6) | +|---|---|---| +| CQL dominates IQN at magnitude head 30-50× (`cql_grad=1-7` vs `iqn_grad=0.05-0.14`) | Pearl 2 per-branch loss budgets in ISV | Budget averaging collapses mag vs dir differentiation | +| NoisyNet σ=0.0226 dominates Q-diff=0.005 by 4-5× | Pearl 3 per-branch σ in ISV, tracks per-branch Q-magnitude | σ averaging re-applies direction's σ to magnitude's narrower Q range | +| `dist_q/h/f` freezes at Bin(2,0.5): 0.225/0.280/0.495 | Pearl 5 per-branch τ in ISV, tracks per-branch Q-skew | τ averaging removes the per-branch concentration benefit | + +Pearl 1 (per-branch C51 atom span) is NOT affected — `atoms_update_kernel.cu` already consumes per-branch ISV at Layer B. +Pearl 4 (per-group Adam hyperparams) is NOT affected — SP4 Layer B's per-group Adam split already dispatches per-group sub-launches. +Pearls 6, 7, 8 are NOT affected — their consumer contracts are either single-direction (Pearl 8) or single-global (Pearl 6) by design. + +SP6 closes the three remaining gaps. + +--- + +## 2. Investigation questions + +These must be answered before SP6 is scoped. Do not start implementation until the L40S smoke answers Q1. + +**Q1 — Does SP5 alone resolve magnitude collapse?** + +Run the L40S smoke at SP5 HEAD (multi_fold_convergence, 5-epoch single-fold). Check: +- Does `var(Q[mag])` grow away from O(1e-9)? +- Does `intent_dist_q` drift away from Bin(2,0.5) over 5 epochs? +- Does `v_half_mag` < `v_half_dir` (Pearl 1 visibly adapting)? + +If yes to all three: SP6 Pearl 2 is still useful (gradient budget) but Pearl 3 and Pearl 5 are lower priority. If intent_dist stays frozen at Bin(2,0.5), Pearl 3 is the likely cause. + +**Q2 — What is the per-branch gradient routing inside grad_buf?** + +`branch_slice_starts_dev` at `gpu_dqn_trainer.rs:13209` already encodes the answer: +- Branch `b` owns tensors `(8 + 4b)..(12 + 4b)` in the padded flat layout. +- Branch 0 (Direction): tensors 8-11 (`w_b0fc`, `b_b0fc`, `w_b0out`, `b_b0out`) +- Branch 1 (Magnitude): tensors 12-15 +- Branch 2 (Order): tensors 16-19 +- Branch 3 (Urgency): tensors 20-23 + +The trunk (tensors 0-7) and value head (tensors 24-33 roughly) sit outside the 4-slice range. The existing `branch_grad_balance_rescale` kernel already operates on these exact slices via `branch_slice_starts_dev` and `branch_slice_lens_dev`. Pearl 2 Approach B can reuse this infrastructure directly. + +But open question: `apply_c51_budget_scale` and `apply_cql_saxpy` currently operate on the FULL `grad_buf` (trunk + value + all branches). If Pearl 2 restricts scaling to only branch slices, the trunk/value head gradient contributions from C51/CQL are no longer budget-scaled. Is that acceptable? Or should Pearl 2 budget-scale trunk separately (shared budget = mean of 4 branch budgets) while applying per-branch budget to branch slices only? + +**Q3 — How many call sites thread noise_sigma, and is ExperienceCollectorConfig re-constructed per step or shared?** + +From `training_loop.rs:1683`, `ExperienceCollectorConfig` is constructed fresh each `collect_experiences_gpu` call. So Pearl 3's change only needs to touch `training_loop.rs:1747` (one call site). But: `ExperienceCollectorConfig.noise_sigma: f32` is a single scalar. To support 4 per-branch values, the struct must grow a `noise_sigma_per_branch: [f32; 4]` field and the kernel at `experience_kernels.cu:3621` must dispatch per-branch. + +How many places construct `ExperienceCollectorConfig`? If there are test harnesses that construct it directly, those are blast radius. + +**Q4 — For Pearl 5: does Pearl 1 per-branch atom span already address the magnitude resolution issue?** + +Pearl 5's motivation is that a symmetric uniform τ schedule is suboptimal for skewed Q-distributions. But if Pearl 1 has already shrunk magnitude's atom span to match its Q-scale (narrow, symmetric around small values), the skew signal may be small and Pearl 5's τ adjustment negligible. This is the strongest argument for deferring Pearl 5 entirely. + +**Q5 — Does the 50-epoch L40S validation (Layer C exit gate) show intent_dist still frozen at Bin(2,0.5) after full SP5?** + +If yes, Pearl 3 is a candidate fix. If no, Pearl 3 may not be needed. Do not run SP6 work until this is known. + +--- + +## 3. Three approaches per consumer + +### Pearl 2 — per-branch loss budgets + +The core obstacle is that `apply_c51_budget_scale`, `apply_cql_saxpy`, and `apply_iqn_trunk_gradient` operate on a flat `grad_buf` with a single scalar multiplier. There is no existing per-branch dispatch. + +**Approach A (minimal — pass per-branch array, scale inside launcher)** + +Extend `compute_adaptive_budgets` to return `([f32;4], [f32;4], [f32;4], [f32;4])`. Extend `apply_c51_budget_scale`, `apply_cql_saxpy`, `apply_iqn_trunk_gradient` to accept `[f32;4]` arrays. Inside each function, loop over branches: compute each branch's start/len from `branch_slice_starts_dev`/`branch_slice_lens_dev`, call a sub-kernel or cuBLAS SCAL on each slice with that branch's scalar. + +The trunk/value portion is scaled by the mean of the 4 branch budgets (preserving current behavior for non-branch parameters). + +- Estimated LOC: ~200 (function signature changes + 4-branch dispatch loops) +- Blast radius: 3 function signatures in `gpu_dqn_trainer.rs`, 1 call site in `fused_training.rs`, `compute_adaptive_budgets` return type +- Perf impact: 4 sub-kernel dispatches per component instead of 1. Overhead negligible vs C51/IQN backward compute. +- Recommendation: preferred starting point — reuses existing `branch_slice_starts_dev` infrastructure, cleanest contract. + +**Approach B (per-branch sub-launches — mirror SP4 Layer B pattern)** + +Identical to Approach A but each sub-launch is a separate CUDA kernel dispatch into the CUDA Graph (if captured) or stream (if not). This is exactly the SP4 Layer B pattern in `launch_sp4_adam_per_group`. + +The difference from Approach A: Approach A loops in Rust host code calling cuBLAS SCAL or a simple scale kernel 4 times. Approach B would make each sub-launch a separate graph node. This matters only if CUDA Graph capture is active. + +- Estimated LOC: ~400 (one per-branch SCAL kernel + graph capture additions for each of 4 sub-launches × 3 components = 12 sub-launches) +- Blast radius: larger — touches graph capture paths in `graph_forward` / `graph_adam` +- Perf impact: marginally better on H100 (graph eliminates launch overhead) but measurable difference is < 0.1% of step time +- Recommendation: only if Approach A shows graph capture errors in smoke + +**Approach C (per-branch scalar axis in the kernel itself)** + +Add a `budget_per_branch[4]` arg to the C51 scale kernel and CQL SAXPY kernel; each thread looks up which branch its weight belongs to from the padded layout and applies that branch's scalar. This requires a per-element branch-membership lookup. + +- Estimated LOC: ~600 (2 new kernel args + kernel body changes + layout encode/decode) +- Blast radius: kernel ABI changes, potentially breaks graph capture if kernel args change +- Perf impact: possible warp divergence (branches in different warps load different scalars). Neutral or slightly negative vs Approach A. +- Recommendation: not recommended. Adds kernel complexity for no benefit over Approach A. + +**Recommendation for Pearl 2: Approach A.** Extend signatures to `[f32;4]`, dispatch 4 cuBLAS SCAL sub-calls per component. Reuse `branch_slice_starts_dev`. + +--- + +### Pearl 3 — per-branch NoisyNet σ + +The obstacle is that `ExperienceCollectorConfig.noise_sigma: f32` is a single scalar and `experience_kernels.cu` uses it globally across all branches. + +**Approach A (minimal — add noise_sigma_per_branch field alongside existing field)** + +Add `noise_sigma_per_branch: [f32; 4]` to `ExperienceCollectorConfig`. In `experience_kernels.cu`, replace the single `noise_sigma` load with a per-branch lookup: `sigma = noise_sigma_per_branch[branch_idx]` (where `branch_idx` is derivable from the action dimension index within the kernel). The existing `noise_sigma: f32` field can be retired once validated. + +`training_loop.rs:1747` is replaced with a 4-element build that reads `ISV[NOISY_SIGMA_BASE..+4]` directly with per-slot cold-start floor. + +- Estimated LOC: ~120 (struct field + 4-element build + kernel dispatch change) +- Blast radius: `ExperienceCollectorConfig` struct (how many test constructors?), `experience_kernels.cu` noise injection site +- Perf impact: 0 (one extra register load per thread to look up branch index) +- Recommendation: preferred. + +**Approach B (per-branch sub-launches of collect_experiences_gpu)** + +Run 4 separate `collect_experiences_gpu` calls, one per branch, each with its branch's σ. Merge results. This is architecturally expensive: `collect_experiences_gpu` is the experience collection kernel itself; 4× calls = 4× experience collection. + +- Estimated LOC: ~50 (trivial call site change) +- Blast radius: high — 4× experience collection latency per step +- Perf impact: 4× experience collection runtime +- Recommendation: do not do this. + +**Approach C (sigma array in kernel, branch-indexed)** + +Equivalent to Approach A in the kernel. The distinction is whether the struct uses `[f32;4]` (Approach A) or a raw pointer to a device buffer (Approach C). Device buffer would require a new GPU allocation; struct field is simpler and CPU-accessible for the host-side construction without violating `feedback_no_htod_htoh_only_mapped_pinned` (it is written into a host-side struct that flows into the kernel as a launch arg, not a DtoH copy). + +- Recommendation: Approach A is already the correct approach; Approach C is Approach A by another name. + +**Recommendation for Pearl 3: Approach A.** One new `[f32;4]` struct field, one kernel change, `training_loop.rs:1747` migrated to per-branch read. + +Note on blast radius investigation: before committing, count `ExperienceCollectorConfig` constructors. The `Default::default()` blanket likely handles most; explicit construction in tests needs checking. + +--- + +### Pearl 5 — per-branch IQN τ schedule + +This is the hardest consumer. The IQN forward pass (`gpu_iqn_head.rs:refresh_taus_from_isv`) produces a single `[B, N]` tau tensor shared across all IQN passes. The architecture does not currently have a per-branch tau axis anywhere downstream. + +**Approach A (minimal — 4 independent IQN forward passes with per-branch τ)** + +After SP6 Pearl 5 Layer B: replace the single `refresh_taus_from_isv` call with 4 calls, one per branch, each setting `online_taus` and `target_taus` to the branch's τ schedule. Then run the IQN forward pass 4 separate times (one per branch), accumulate per-branch distributional Q estimates. Merge via existing `q_quantile_reduce` or a new per-branch reduce. + +- Estimated LOC: ~400 (4 tau arrays, 4 IQN forward passes, merge kernel or loop) +- Blast radius: IQN forward call graph, `GpuIqnHead` struct (needs to hold 4 tau tensors), `fused_training.rs` IQN dispatch loop +- Perf impact: 3-4× IQN forward compute per step. IQN is a significant fraction of step time on L40S. Rough estimate: +15-25% total step time. +- Recommendation: viable for correctness validation, but expensive. + +**Approach B (per-branch tau axis in the IQN kernel)** + +Add a `branch_idx` axis to the tau tensor: `[B, 4, N]` instead of `[B, N]`. Each IQN forward uses `tau_tensor[:, branch_idx, :]` for its branch. The kernel would need a new argument and an additional stride. + +- Estimated LOC: ~800 (kernel arg + stride change + all IQN forward call sites + tau tensor alloc growth) +- Blast radius: `GpuIqnHead`, `gpu_iqn_head.rs`, `fused_training.rs`, any test that constructs tau tensors +- Perf impact: neutral (same number of ops; marginally better memory access pattern if branches are contiguous) +- Recommendation: best perf/complexity tradeoff if Pearl 5 is included in SP6. + +**Approach C (do nothing — defer Pearl 5 indefinitely)** + +Pearl 5's benefit is τ schedule adaptation for per-branch Q-skew. If Pearl 1 (per-branch atom span) has already correct-sized the C51 distribution to each branch's Q-scale, the Q-skew within each branch's narrow distribution is small. The τ schedule adjustment (`skew × SKEW_SHIFT=0.05`, max ±0.15) may produce negligible improvement given well-adapted atoms. + +The `refresh_taus_from_isv` averaging (mean of 4 branches per quantile) produces exactly the result Pearl 5 would produce if all branches had identical Q-distributions — which is approximately the case for uniform-policy cold-start and may remain approximately true for Order/Urgency branches. + +- Estimated LOC: 0 +- Blast radius: 0 +- Perf impact: 0 +- Recommendation: strongly worth considering. Defer until Q4 below is answered. + +**Open decision: include Pearl 5 in SP6 or defer?** + +This cannot be resolved without first answering Q4 (does Pearl 1 make Pearl 5 negligible?) and Q5 (does intent_dist still freeze post-SP5?). Pearl 5's infrastructure work (Approach B) is substantial and the expected benefit is the smallest of the three consumers. Deferral is the default recommendation unless Q4+Q5 show it is necessary. + +--- + +## 4. Sequencing + +**Independence:** The three consumers are fully independent from each other. + +- Pearl 2 touches `grad_buf` scaling (backward pass). +- Pearl 3 touches experience collection (forward pass). +- Pearl 5 touches IQN tau tensors (distributional Q forward). + +There are no shared buffers, no ordering constraints, no circular dependencies between them. + +**Commit granularity:** Per `feedback_no_partial_refactor`, each pearl is a single atomic commit. Consumer and any companion test land together. No partial migration (e.g., extending the struct field without migrating the kernel call site) is acceptable. + +**Suggested order if all three proceed:** + +1. Pearl 2 first — directly addresses the CQL-dominates-IQN finding, highest expected impact, lowest kernel complexity. +2. Pearl 3 second — medium impact (σ-to-Q-diff ratio), low kernel complexity, requires blast radius audit of `ExperienceCollectorConfig`. +3. Pearl 5 last — highest kernel complexity, lowest expected incremental gain post-Pearl-2+3. Omit entirely if Q4+Q5 answers permit. + +**Layer structure:** + +SP6 does not have a "Layer A" (no new producers; all producers are SP5 Layer A). It is pure "Layer B" work: consumer migration commits only. Each commit is independently buildable and reversible. No Layer A / Layer D equivalent. + +**Validation per commit:** + +- `cargo check -p ml --lib` clean +- `multi_fold_convergence` smoke passes (local RTX 3050 Ti, 5-epoch single-fold) +- HEALTH_DIAG confirms per-branch budget differentiation visible post-Pearl-2 (e.g., `last_cql_budget_eff` should now vary per-branch; need a new diagnostic field or log line) + +--- + +## 5. Alternative: Pearl 2 only + +Pearl 2 is the unique structural gap. The original finding was a 30-50× CQL-vs-IQN imbalance at the magnitude head specifically. With averaging, CQL's budget at the magnitude head is diluted by direction/order/urgency budgets that are presumably less imbalanced — so CQL still over-dominates IQN at magnitude, just less visibly. + +The case for "Pearl 2 only" as the full scope of SP6: + +- Pearl 1 (per-branch atom span) directly fixes the magnitude Q-resolution mismatch that was the root cause of both the σ-to-Q-diff ratio (Pearl 3's motivation) and the τ-skew mismatch (Pearl 5's motivation). If Pearl 1 successfully shrinks magnitude's atom span from [-15,15] to [-0.025,0.025], then: + - The new σ produced by Pearl 3's ISV producer is already ~100× smaller (proportional to `v_half`). Even averaged with direction's larger σ, the mean is smaller than the pre-SP5 global σ. The σ-to-Q-diff pathology is reduced by Pearl 1 alone. + - Pearl 5's τ schedule assumes the distribution is skewed. Once Pearl 1 properly centers and scales the magnitude distribution, skew is minimal. τ schedule adaptation is marginal. +- Pearl 2's per-branch budget directly corrects the CQL-vs-IQN loss component imbalance AT the magnitude head — which is a budget problem independent of atom span. + +Therefore: the minimal viable SP6 is Pearl 2 only. Pearl 3 and Pearl 5 are "nice to have" if Pearl 1+2 leave residual issues. + +**Condition under which Pearl 3 becomes required:** L40S 50-epoch validation post-SP5+Pearl2 still shows intent_dist frozen at Bin(2,0.5) or σ/Q-diff ratio > 2× at magnitude head. + +**Condition under which Pearl 5 becomes required:** Same validation shows IQN τ mismatch visible in distributional Q-skew diagnostics (Pearl 5's producer signal `ISV[SKEW_KURTOSIS_BASE]` consistently nonzero at magnitude head, and the averaged τ schedule visibly misaligned). + +--- + +## 6. Investigation gates before SP6 starts + +SP6 should not begin implementation until all three gates pass: + +**Gate 1 — L40S smoke at SP5 HEAD** + +Run `multi_fold_convergence` (5-epoch, single-fold, L40S) at SP5 HEAD (`3ad5e011b`). Kill at first HEALTH_DIAG line per `feedback_kill_runs_on_anomaly_quickly`. Check: +- `v_half_mag` vs `v_half_dir` visible separation +- `intent_dist_q/h/f` drifting vs frozen at Bin(2,0.5) +- `var_q_mag` growing vs flat O(1e-9) + +**Gate 2 — 50-epoch validation** + +Run `train-multi-seed-baseline` (3-seed × 50-epoch × 3-fold, L40S). SP5 Layer C exit gate. If `intent_dist` resolves at this validation, Pearl 3 and Pearl 5 may be unnecessary. + +**Gate 3 — budget differentiation diagnostic** + +To know whether Pearl 2 is working post-implementation, the diagnostic infrastructure must be extended. Currently `last_iqn_budget_eff` / `last_cql_budget_eff` (`fused_training.rs:3327`) are single scalars. Post-Pearl-2, 4-element arrays are needed. HEALTH_DIAG (`health_diag.rs`) needs a `cql_budget_per_branch: [f32;4]` field (or similar) so training logs expose the differentiation. This is not a large change but must be included in Pearl 2's commit. + +--- + +## 7. Out of scope for SP6 + +- **Pearl 1 per-branch atom span** — already complete in SP5 Layer B +- **Pearl 4 per-group Adam hyperparams** — already complete in SP5 Layer B (per-group sub-launches) +- **Pearl 6 Kelly cap signals** — single-domain consumer, no per-branch issue +- **Pearl 8 trail-stop per-direction** — single-direction consumer, no averaging issue +- **HEALTH_DIAG GPU port** — separate deferred project; current HEALTH_DIAG is acceptable for SP6 validation if log fields are extended for per-branch budget visibility +- **TFT backward pipeline** — deferred separate project +- **Layer D host-EMA → GPU close-out** — SP5 deferred separate atomic commit, not SP6 scope +- **Pearl 7 `intent_dist` Bin(2,0.5) investigation** — SP5 Layer C step 4; resolved or not by 50-epoch validation. If not resolved, a separate fix, not SP6. + +--- + +## Open decisions + +**D1.** Does the 5-epoch L40S smoke at SP5 HEAD show `v_half_mag` < `v_half_dir`? (Resolves whether Pearl 1 is actually adapting, which changes the Pearl 3 + Pearl 5 priority.) + +**D2.** Does the 50-epoch validation show `intent_dist` still frozen at Bin(2,0.5)? (Resolves Pearl 3 necessity.) + +**D3.** For Pearl 2 Approach A: should `apply_c51_budget_scale` / `apply_cql_saxpy` scale the trunk/value region using the mean of 4 branch budgets (preserving current trunk behavior) or using a fixed scalar of 1.0 (no trunk budget-scaling)? The trunk is shared — it seems reasonable to scale trunk by the mean. But the SP5 design spec is silent on this detail. + +**D4.** Pearl 5 — include or defer? Default: defer pending D1+D2 answers. + +**D5.** What is the blast radius of extending `ExperienceCollectorConfig` to `noise_sigma_per_branch: [f32;4]`? Are there test-only constructors that hardcode the struct? (Must be audited before Pearl 3 commit.) + +--- + +End of brainstorm. No implementation decisions made; all five open decisions require validation data before proceeding.