docs(sp6): brainstorm spec — per-branch consumer infrastructure refactor

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

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.