# d_adv_logits Branch-Major Layout 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 the d_adv_logits memory layout mismatch that causes zero gradients at H100 scale (batch=16384). The gradient kernels write sample-major interleaved `[B, total_branch_atoms]` but the cuBLAS backward reads branch-major `[B*B0*NA | B*B1*NA | B*B2*NA | B*B3*NA]`.

**Architecture:** Change the atomicAdd indexing formula in the 3 CUDA gradient kernels (c51_grad, mse_grad, cql_grad) to write in branch-major layout. The Rust allocations, blend SAXPY, f32→bf16 cast, and cuBLAS backward pointer arithmetic are all layout-agnostic or already branch-major — they don't change.

**Tech Stack:** CUDA kernels (.cu), Rust orchestrator, cuBLAS GemmEx backward.

---

### Current layout (BROKEN at scale)

```
d_adv_logits[b * total_branch_atoms + branch_offset + a * NA + j]
```
Where `branch_offset = sum_{d'<d}(branch_sizes[d'] * NA)` — offset within a SINGLE sample's branch atoms. This interleaves all 4 branches per sample: `[S0_B0, S0_B1, S0_B2, S0_B3, S1_B0, S1_B1, ...]`

### Target layout (matches cuBLAS backward)

```
d_adv_logits[branch_base + b * (A_d * NA) + a * NA + j]
```
Where `branch_base = sum_{d'<d}(B * branch_sizes[d'] * NA)` — offset to the START of branch d's contiguous block across ALL samples: `[ALL_S_B0, ALL_S_B1, ALL_S_B2, ALL_S_B3]`

### Why this works

The cuBLAS backward pointer arithmetic at `gpu_dqn_trainer.rs:5540-5543` already computes:
```rust
d_adv1 = d_adv0 + (B * b0 * na) * bf16_size
```
This is branch-major. After the fix, the gradient kernels write in the same layout, so the flat f32→bf16 cast preserves the layout, and cuBLAS reads correct per-branch `[B, n_d]` matrices.

---

### Task 1: Fix c51_grad_kernel.cu — branch-major d_adv_logits indexing

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/c51_grad_kernel.cu:75-97`

- [ ] **Step 1: Change the branch_offset and adv_idx computation**

In `c51_grad_kernel.cu`, find (around line 81-97):

```cuda
    int branch_offset = 0;
    for (int dd = 0; dd < d; dd++)
        branch_offset += branch_sizes[dd] * num_atoms;

    /* ... */

    /* d_adv[b, d, a, j] = branch_scale * d_combined * (delta(a, a_d) - 1/A_d) */
    for (int a = 0; a < A_d; a++) {
        float dueling_grad = (a == a_d) ? (1.0f - inv_A) : (-inv_A);
        float grad_val = branch_scale * d_combined * dueling_grad;
        int adv_idx = b * total_branch_atoms + branch_offset + a * num_atoms + j;
        atomicAdd(&d_adv_logits[adv_idx], grad_val);
    }
```

Replace with:

```cuda
    /* Branch-major base offset: all B samples for branches 0..d-1 precede branch d.
     * Layout: [B*B0*NA | B*B1*NA | B*B2*NA | B*B3*NA] (contiguous per-branch blocks).
     * Matches cuBLAS backward's pointer arithmetic for per-branch dY matrices. */
    int branch_base = 0;
    for (int dd = 0; dd < d; dd++)
        branch_base += batch_size * branch_sizes[dd] * num_atoms;

    /* ... */

    /* d_adv[branch_base + b * (A_d * NA) + a * NA + j] — branch-major layout */
    for (int a = 0; a < A_d; a++) {
        float dueling_grad = (a == a_d) ? (1.0f - inv_A) : (-inv_A);
        float grad_val = branch_scale * d_combined * dueling_grad;
        int adv_idx = branch_base + b * (A_d * num_atoms) + a * num_atoms + j;
        atomicAdd(&d_adv_logits[adv_idx], grad_val);
    }
```

- [ ] **Step 2: Build**

Run: `SQLX_OFFLINE=true cargo check -p ml`
Expected: compiles clean (the .cu is compiled by build.rs via nvcc).

- [ ] **Step 3: Commit**

```
git add crates/ml/src/cuda_pipeline/c51_grad_kernel.cu
git commit -m "fix(critical): c51_grad writes branch-major d_adv_logits layout

The old sample-major interleaved layout [B, total_branch_atoms] did not
match cuBLAS backward's branch-major pointer arithmetic [B*B0*NA, B*B1*NA, ...].
At batch=16384, the backward GEMM read cross-contaminated data from wrong
branches, producing zero weight gradients by symmetry cancellation."
```

---

### Task 2: Fix mse_grad_kernel.cu — same branch-major indexing

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/mse_grad_kernel.cu:80-97`

- [ ] **Step 1: Change the branch_offset and adv_idx computation**

In `mse_grad_kernel.cu`, find (around line 80-97):

```cuda
    int branch_offset = 0;
    for (int dd = 0; dd < d; dd++)
        branch_offset += branch_sizes[dd] * num_atoms;

    /* ... */

    /* d_adv[b, d, a, j] = branch_scale * d_combined * (delta(a, a_d) - 1/A_d) */
    for (int a = 0; a < A_d; a++) {
        float dueling_grad = (a == a_d) ? (1.0f - inv_A) : (-inv_A);
        float grad_val = branch_scale * d_combined * dueling_grad;
        int adv_idx = b * total_branch_atoms + branch_offset + a * num_atoms + j;
        atomicAdd(&d_adv_logits[adv_idx], grad_val);
    }
```

Replace with:

```cuda
    /* Branch-major layout — see c51_grad_kernel.cu for full explanation. */
    int branch_base = 0;
    for (int dd = 0; dd < d; dd++)
        branch_base += batch_size * branch_sizes[dd] * num_atoms;

    /* ... */

    /* d_adv[branch_base + b * (A_d * NA) + a * NA + j] */
    for (int a = 0; a < A_d; a++) {
        float dueling_grad = (a == a_d) ? (1.0f - inv_A) : (-inv_A);
        float grad_val = branch_scale * d_combined * dueling_grad;
        int adv_idx = branch_base + b * (A_d * num_atoms) + a * num_atoms + j;
        atomicAdd(&d_adv_logits[adv_idx], grad_val);
    }
```

- [ ] **Step 2: Build**

Run: `SQLX_OFFLINE=true cargo check -p ml`

- [ ] **Step 3: Commit**

```
git add crates/ml/src/cuda_pipeline/mse_grad_kernel.cu
git commit -m "fix(critical): mse_grad writes branch-major d_adv_logits layout"
```

---

### Task 3: Fix cql_grad_kernel.cu — branch-major d_adv_logits indexing

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/cql_grad_kernel.cu:140-190`

The CQL kernel writes per-sample output (`d_adv_logits[i, ...]`), not via atomicAdd across samples. But it still uses the sample-major interleaved layout. The CQL backward in `apply_cql_gradient()` then reads this buffer and backward-passes it through cuBLAS with the same branch-major pointer arithmetic.

- [ ] **Step 1: Change the CQL adv_logits indexing to branch-major**

In `cql_grad_kernel.cu`, find the inner loop (around line 140-193). The CQL kernel iterates `for d in 0..4`, and for each branch, offsets into `d_adv_logits` via:

```cuda
float* d_adv = d_adv_logits + (long long)i * total_actions * num_atoms
             + (long long)(adv_offset + a) * num_atoms;
```

Where `adv_offset = sum_{d'<d}(branch_sizes[d'])` and `total_actions = sum_d(branch_sizes[d])`.

This is sample-major. Change to branch-major:

```cuda
/* Branch-major base: all N samples for branches 0..d-1 precede branch d. */
int branch_base = 0;
for (int dd = 0; dd < d; dd++)
    branch_base += N * branch_sizes[dd] * num_atoms;

/* Per-sample offset within branch d's block */
float* d_adv = d_adv_logits + branch_base + (long long)i * bd * num_atoms
             + (long long)a * num_atoms;
```

Apply this in BOTH places the kernel writes to `d_adv` (the gradient write loop at ~line 186 and the accumulated `d_val_accum` write).

Also update the `d_adv` READS in the forward Q-value computation (lines 77-119) to use branch-major layout from `adv_logits` (the FORWARD logit buffer). **BUT WAIT**: the FORWARD `adv_logits` buffer uses the same sample-major layout. Only the GRADIENT output `d_adv_logits` needs branch-major.

The FORWARD buffer `adv_logits` (read-only in CQL) is a SEPARATE buffer from `d_adv_logits` (write output). The forward buffer's layout is set by `launch_cublas_forward()` which writes branch-major already (each branch GEMM writes to a contiguous `[B, n_d]` region). So the FORWARD reads are correct as-is.

Only change the WRITES to `d_adv_logits`.

- [ ] **Step 2: Build**

Run: `SQLX_OFFLINE=true cargo check -p ml`

- [ ] **Step 3: Commit**

```
git add crates/ml/src/cuda_pipeline/cql_grad_kernel.cu
git commit -m "fix(critical): cql_grad writes branch-major d_adv_logits layout"
```

---

### Task 4: Revert diagnostic hacks, run smoke tests, deploy

**Files:**
- Modify: `crates/ml/src/trainers/dqn/fused_training.rs` — revert `gradient_stage_diagnostics: true` back to env var
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` — revert ungraphed `submit_forward_ops_main()` back to graph capture+replay

- [ ] **Step 1: Revert diagnostic hacks**

In `fused_training.rs`, change:
```rust
gradient_stage_diagnostics: true, // DIAG: force ungraphed path ...
```
Back to:
```rust
gradient_stage_diagnostics: std::env::var("FOXHUNT_GRAD_DIAG").is_ok(),
```

In `gpu_dqn_trainer.rs`, in `train_step_gpu()`, change:
```rust
        // DIAG: skip ALL graph capture — run forward+backward ungraphed every step.
        self.update_stochastic_depth_mask()?;
        self.submit_forward_ops_main()?;
```
Back to:
```rust
        if self.graph_forward.is_none() {
            self.capture_training_graphs(online_dueling, online_branching)?;
        }
        self.update_stochastic_depth_mask()?;
        self.replay_forward()?;
```

- [ ] **Step 2: Build and run full smoke test suite**

Run: `SQLX_OFFLINE=true cargo check --workspace`
Run: `SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests::training_stability --ignored`
Expected: 6/6 pass.

Run: `SQLX_OFFLINE=true cargo test -p ml-core --lib`
Expected: 300/300 pass.

- [ ] **Step 3: Commit**

```
git add -A
git commit -m "fix(critical): branch-major d_adv_logits layout + revert diagnostic hacks

Root cause of H100 gradient collapse: c51_grad, mse_grad, and cql_grad
kernels wrote d_adv_logits in sample-major interleaved layout [B, TBA]
but cuBLAS backward read it as branch-major [B*B0*NA | B*B1*NA | ...].
At batch=16384, the GEMM read cross-contaminated data from wrong branches,
producing zero weight gradients by symmetry cancellation at scale.

Fix: change indexing formula in all 3 grad kernels to write branch-major.
No Rust changes needed — allocations, blend, cast are layout-agnostic."
```

- [ ] **Step 4: Push and deploy H100 baseline**

```
git push
./scripts/argo-train.sh dqn --baseline --epochs 20
```

Expected: grad_norm stable (0.5-1.0) across all 20 epochs on H100.

---

### Task 5: Fix blend n_val/n_adv padding consistency (secondary)

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs:4568-4569`

The blend uses `n_val = b * na` but the cast uses `n_val = b * pad32(na)`. The padding elements are zero so this is harmless, but inconsistent.

- [ ] **Step 1: Align blend dimensions with cast**

Change line 4568-4569 from:
```rust
let n_val = (b * na) as i32;
let n_adv = (b * (b0 + b1 + b2 + b3) * na) as i32;
```
To:
```rust
let n_val = (b * pad32(na)) as i32;
let n_adv = (b * (b0 + b1 + b2 + b3) * na + 32 * 4) as i32;
```

Also update `val_blocks` and `adv_blocks` to match.

- [ ] **Step 2: Build and test**

Run: `SQLX_OFFLINE=true cargo check -p ml`
Run: `SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests::regression --ignored`

- [ ] **Step 3: Commit**

```
git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
git commit -m "fix: align blend n_val/n_adv with cast dimensions for consistency"
```
