feat(sp20): Phase 1.1 sp20_stats_compute kernel
Component 5 / Kernel 3 of the SP20 fused-producer chain. Single-block BLOCK=256 kernel reads `aux_logits [B, 3]` (the SP14-C aux head's 3-class direction logits) and emits `[aux_conf_p50, aux_conf_std]` into a `MappedF32Buffer<2>`, where the per-row signal is `aux_conf[i] = max_c softmax(logits[i, *])[c] - 1/3`. p50 uses the inlined `sp4_histogram_p99` pattern (per-warp tile binning + cumulative-from-bottom, no atomicAdd per `feedback_no_atomicadd`); std uses two block tree-reductions sharing one shmem tile sequentially. One fused kernel streams `aux_logits` once for both stats per `pearl_fused_per_group_statistics_oracle`. Phase 1.4 wires the production launch site atomically with the rest of the SP20 reward chain per `feedback_no_partial_refactor`. This commit lands kernel + Rust launcher + GPU oracle tests + build entry + audit-doc entry together so the kernel is independently verifiable on RTX 3050 Ti (sm_86) and L40S (sm_89) before the EMA + controller producers (Phase 1.2 + 1.3) reference its outputs. Tests verify: - uniform logits → aux_conf = 0 → [p50, std] = [0, 0] - varied confidence (logit ramp 0 → 3) → matches CPU oracle - heterogeneous half-hot half-uniform → matches CPU oracle - empty batch → degenerate-guard writes [0, 0] All 4 GPU oracle tests + 4 launcher unit tests pass on RTX 3050 Ti. Test data uses per-row variance to avoid the `pearl_sp4_histogram_warp_tile_undercount` lockstep-uniform trap (concentrated values within one bin_width race the per-warp non-atomic increments). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -10751,3 +10751,159 @@ bash scripts/audit_sp18_consumers.sh --check # exit
|
||||
round-trip equality. Direction sometimes flips at low Q values,
|
||||
unrelated to ISV slot layout. Verified flaky pre-Commit B by
|
||||
stashing the SP19 changes and re-running 3× — fails 2/3 runs.
|
||||
|
||||
## 2026-05-09 — SP20 Phase 1.1: sp20_stats_compute kernel (additive)
|
||||
|
||||
### Component
|
||||
|
||||
- New CUDA kernel `crates/ml/src/cuda_pipeline/sp20_stats_compute_kernel.cu`
|
||||
- New Rust launcher `crates/ml/src/cuda_pipeline/sp20_stats_compute.rs`
|
||||
(re-exported from `cuda_pipeline::sp20_stats_compute`)
|
||||
- New GPU oracle test `crates/ml/tests/sp20_stats_compute_test.rs`
|
||||
- Cubin manifest entry in `crates/ml/build.rs`
|
||||
|
||||
### Purpose
|
||||
|
||||
Component 5 / Kernel 3 of the SP20 fused-producer chain (Phase 1.4
|
||||
will land Kernels 1+2 — `sp20_emas_compute` + `sp20_controllers_
|
||||
compute` — and the production launch site atomically with the
|
||||
training-loop wire-up per `feedback_no_partial_refactor`).
|
||||
|
||||
Reads `aux_logits [B, 3]` (the SP14-C aux head's 3-class direction
|
||||
logits) and emits `[aux_conf_p50, aux_conf_std]` into a
|
||||
`MappedF32Buffer<2>`, where the per-row signal is
|
||||
|
||||
```
|
||||
aux_conf[i] = max_c softmax(logits[i, *])[c] - 1/3
|
||||
```
|
||||
|
||||
`aux_conf` measures peak class confidence above the K=3 uniform
|
||||
baseline (0 ⇔ uniform, 2/3 ⇔ fully concentrated). The downstream
|
||||
EMA producer (Phase 1.2) Wiener-blends p50 + std into ISV slots
|
||||
within [510..520) reserved by SP20 (`f5eed1fa7`).
|
||||
|
||||
### Algorithm
|
||||
|
||||
Single-block, 256-thread kernel; one fused stream of `aux_logits`
|
||||
for both stats per `pearl_fused_per_group_statistics_oracle`.
|
||||
|
||||
- **Pass A**: per-row numerically-stable softmax → max → minus 1/3,
|
||||
write `aux_conf[i]` into per-row scratch in dynamic shmem.
|
||||
- **Pass B1**: block tree-reduce `sum(aux_conf)`. Reuses
|
||||
`s_bins[0..256]` (the histogram tile) as a float-via-int reduce
|
||||
tile, BLOCK=256 ≤ 256 bins so the alias fits exactly.
|
||||
- **Pass B2**: block tree-reduce `sum(aux_conf²)` using the same
|
||||
shmem tile sequentially.
|
||||
- **Pass C.1**: max-reduce `|aux_conf|` to seed `step_max` for the
|
||||
histogram (mirrors `sp4_histogram_p99` pass 1 verbatim).
|
||||
- **Pass C.2**: per-warp tile binning into 256 linear bins (no
|
||||
atomicAdd per `feedback_no_atomicadd`; mirrors
|
||||
`sp4_histogram_p99` pass 2).
|
||||
- **Pass C.3**: cumulative-from-BOTTOM scan (`cumul ≥ ⌈B/2⌉`) →
|
||||
p50 = bin upper-edge. Inverts the direction of
|
||||
`sp4_histogram_p99` pass 3 (which scans top-down for the 99th
|
||||
percentile); the scan logic itself is otherwise identical.
|
||||
|
||||
Thread 0 writes `out[0] = p50`, `out[1] = sqrt(var)` with
|
||||
`__threadfence_system()` for mapped-pinned coherence.
|
||||
|
||||
### Wiring contract
|
||||
|
||||
Phase 1.1 wires kernel + launcher + tests + build entry **only**.
|
||||
There is NO production caller in this commit — the kernel is dead
|
||||
code awaiting Phase 1.4's atomic wire-up of:
|
||||
|
||||
1. The training-loop call site (post-aux-head-forward, pre-EMA
|
||||
producer) supplying the `aux_logits` device pointer.
|
||||
2. A trainer-owned `MappedF32Buffer<2>` for the `[p50, std]` output.
|
||||
3. Stream ordering with the aux-head forward (same-stream is
|
||||
sufficient; CUDA stream-ordering invariant guarantees the
|
||||
`aux_logits` write is visible before this kernel reads it).
|
||||
|
||||
The `cuda_pipeline::sp20_stats_compute` module exports
|
||||
`launch_sp20_stats_compute` + the contract constants (`AUX_K_CLASSES`,
|
||||
`SP20_STATS_BLOCK`, `SP20_STATS_HIST_BINS`, `dynamic_shmem_bytes`)
|
||||
the wire-up will consume.
|
||||
|
||||
### Pearls + invariants honoured
|
||||
|
||||
- `feedback_no_atomicadd` — block tree-reduce + per-warp tile
|
||||
binning; no atomicAdd anywhere.
|
||||
- `feedback_no_cpu_compute_strict` — every operation GPU-side.
|
||||
- `feedback_no_htod_htoh_only_mapped_pinned` — output is a
|
||||
`MappedF32Buffer<2>`; kernel emits `__threadfence_system()`.
|
||||
- `feedback_no_partial_refactor` — kernel + launcher + tests +
|
||||
build entry land in one commit; production wire-up lands
|
||||
atomically in Phase 1.4 with the EMA + controller producers.
|
||||
- `pearl_fused_per_group_statistics_oracle` — one kernel
|
||||
computing both p50 and std from a single stream of `aux_logits`.
|
||||
- `pearl_no_host_branches_in_captured_graph` — single-block
|
||||
kernel, captureable.
|
||||
- `pearl_first_observation_bootstrap` — degenerate-distribution
|
||||
branch (step_max == 0) writes `[0, 0]` (sentinel) so the
|
||||
downstream EMA producer can apply Pearl-A bootstrap on the
|
||||
next non-degenerate observation.
|
||||
- `pearl_sp4_histogram_warp_tile_undercount` (test data) — test
|
||||
vectors use varied per-row confidences (logit ramps, jittered
|
||||
half-hot/half-uniform) to avoid the lockstep-uniform trap that
|
||||
causes warp-tile race undercounts in the non-atomic per-warp
|
||||
histogram.
|
||||
|
||||
### Tests
|
||||
|
||||
`crates/ml/tests/sp20_stats_compute_test.rs` — GPU oracle tests
|
||||
(all `#[ignore = "requires GPU"]`):
|
||||
|
||||
1. `uniform_logits_emit_zero_stats` — `aux_logits = 0` ⇒
|
||||
softmax = [1/3, 1/3, 1/3] ⇒ aux_conf = 0 ⇒ p50 = std = 0.
|
||||
Exercises the kernel's degenerate-distribution branch.
|
||||
2. `varied_confidence_logits_match_cpu_oracle` — class-0 logit
|
||||
ramps 0 → 3, others 0; CPU oracle matches GPU within
|
||||
bin_width tolerance for p50 (≈ 0.002) and 5e-3 for std.
|
||||
3. `heterogeneous_logits_match_cpu_oracle` — half jittered
|
||||
hot, half jittered uniform; CPU oracle matches GPU within
|
||||
0.02 for p50 and 5e-3 for std.
|
||||
4. `empty_batch_writes_zero_stats` — `B = 0` exercises the
|
||||
pre-warmup degenerate guard; kernel writes `[0, 0]` and
|
||||
returns without entering the histogram path.
|
||||
|
||||
Plus 4 unit tests in the launcher module (`AUX_K_CLASSES = 3`,
|
||||
block-size discipline, shmem budget under L40S 48 KiB, zero-batch
|
||||
shmem accounting).
|
||||
|
||||
### Files modified
|
||||
|
||||
| File | Status | Purpose |
|
||||
|------|--------|---------|
|
||||
| `crates/ml/src/cuda_pipeline/sp20_stats_compute_kernel.cu` | NEW | Single-block stats producer kernel |
|
||||
| `crates/ml/src/cuda_pipeline/sp20_stats_compute.rs` | NEW | Rust launcher + contract constants |
|
||||
| `crates/ml/tests/sp20_stats_compute_test.rs` | NEW | 4 GPU oracle tests |
|
||||
| `crates/ml/src/cuda_pipeline/mod.rs` | +pub mod | Module declaration |
|
||||
| `crates/ml/build.rs` | +cubin entry | Compile entry for nvcc |
|
||||
| `docs/dqn-wire-up-audit.md` | This entry | Audit log |
|
||||
|
||||
### Verification
|
||||
|
||||
```
|
||||
SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 cargo check -p ml --features cuda
|
||||
SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 cargo test -p ml \
|
||||
--test sp20_stats_compute_test --features cuda -- --ignored --nocapture
|
||||
SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 cargo test -p ml \
|
||||
--features cuda --lib sp20_stats_compute
|
||||
```
|
||||
|
||||
All 4 GPU oracle tests + 4 unit tests pass on RTX 3050 Ti (sm_86).
|
||||
L40S verification deferred until Phase 1.4 lands the production
|
||||
caller.
|
||||
|
||||
### Phase 1.1 → Phase 1.4 forward references
|
||||
|
||||
- Phase 1.2: `sp20_emas_compute_kernel.cu` consumes
|
||||
`[aux_conf_p50, aux_conf_std]` (this kernel's output) +
|
||||
per-trade `R_event` / WR-window observations to advance the
|
||||
Wiener EMAs in ISV slots [510..520).
|
||||
- Phase 1.3: `sp20_controllers_compute_kernel.cu` derives
|
||||
`LOSS_CAP / TARGET_HOLD_PCT / N_STEP / AUX_CONF_THRESHOLD /
|
||||
AUX_GATE_TEMP / HOLD_COST_SCALE` from the EMAs.
|
||||
- Phase 1.4: training-loop wire-up calls all 3 producers in
|
||||
sequence on the same stream as the aux-head forward.
|
||||
|
||||
Reference in New Issue
Block a user