diff --git a/docs/superpowers/plans/2026-04-27-moe-regime-redesign.md b/docs/superpowers/plans/2026-04-27-moe-regime-redesign.md index 82de7d69c..fce2b1d88 100644 --- a/docs/superpowers/plans/2026-04-27-moe-regime-redesign.md +++ b/docs/superpowers/plans/2026-04-27-moe-regime-redesign.md @@ -12,6 +12,8 @@ **Pearl:** [`pearl_learned_gate_subsumes_handcoded.md`](../../../.claude/projects/-home-jgrusewski-Work-foxhunt/memory/pearl_learned_gate_subsumes_handcoded.md). Load-bearing rationale. +**Memory-transfer constraint (`feedback_no_htod_htoh_only_mapped_pinned.md`):** Never use HtoD (`memcpy_stod`) or HtoH (defensive Vec→Vec copies) anywhere in this plan — including test wrappers. The only allowed CPU↔GPU path is **mapped pinned memory** via `cuMemHostAlloc(DEVICEMAP | PORTABLE)`. Reference implementation: `MappedF32Buffer` / `MappedI32Buffer` in `crates/ml/src/trainers/dqn/distributional_q_tests.rs` (used by Test 0.F). Every kernel input that originates CPU-side allocates a mapped buffer, the CPU writes to `host_ptr`, the kernel reads via `dev_ptr` — no copy step. Every kernel output that the CPU needs to read is similarly written by the kernel via `dev_ptr` and read CPU-side via `host_ptr` (with `__threadfence_system()` on the kernel side and `read_volatile` on the CPU side). + --- ## File Structure @@ -796,7 +798,91 @@ Co-Authored-By: Claude Opus 4.7 (1M context) " ## Phase 2 — MoE kernels -Phase 2 implements the 4 new CUDA kernels with unit tests. Kernels are written and tested in isolation; not yet wired into training graph. +Phase 2 implements the 4 new CUDA kernels with unit tests. Kernels are written and tested in isolation; not yet wired into training graph. **All test wrappers use mapped pinned memory exclusively** per `feedback_no_htod_htoh_only_mapped_pinned.md` — no `memcpy_stod`, no `memcpy_dtov`, no defensive Vec-to-Vec. + +### Task 2.0 — Promote MappedF32Buffer / MappedI32Buffer to shared module + +**Files:** +- Create: `crates/ml/src/cuda_pipeline/mapped_pinned.rs` +- Modify: `crates/ml/src/cuda_pipeline/mod.rs` (add `pub mod mapped_pinned;`) +- Modify: `crates/ml/src/trainers/dqn/distributional_q_tests.rs` (delete local definitions, import from shared) + +The existing Test 0.F (commit `da632446c`) defines `MappedF32Buffer` and `MappedI32Buffer` locally in `distributional_q_tests.rs`. They use `cuMemHostAlloc(DEVICEMAP | PORTABLE)` + `cuMemHostGetDevicePointer_v2` — exactly the pattern this plan needs everywhere. Promote to a shared module so all kernel-test wrappers share one implementation; delete the local copies. + +- [ ] **Step 1: Locate the existing buffer definitions** + +```bash +grep -n "struct MappedF32Buffer\|struct MappedI32Buffer\|impl Mapped" crates/ml/src/trainers/dqn/distributional_q_tests.rs | head -10 +``` + +- [ ] **Step 2: Create the shared module** + +Create `crates/ml/src/cuda_pipeline/mapped_pinned.rs` by moving the entire `MappedF32Buffer` and `MappedI32Buffer` blocks (struct + `unsafe impl Send/Sync` + `impl new` + `impl read_all` + `impl Drop`) from `distributional_q_tests.rs` into the new file. Make the structs and methods `pub`. + +Add a helper method per buffer type for "write from slice without memcpy": + +```rust +impl MappedF32Buffer { + /// Write CPU-side data via host_ptr — direct memory write, no memcpy. + /// Caller must ensure `slice.len() <= self.len`. + pub fn write_from_slice(&self, slice: &[f32]) { + assert!(slice.len() <= self.len, "write_from_slice overflow"); + // Safety: host_ptr is valid for self.len f32s; writes are coherent + // with subsequent kernel reads via dev_ptr (cuMemHostAlloc DEVICEMAP). + unsafe { + for (i, &v) in slice.iter().enumerate() { + std::ptr::write_volatile(self.host_ptr.add(i), v); + } + } + } +} +``` + +(Same for `MappedI32Buffer` with `i32`.) + +- [ ] **Step 3: Add module to mod.rs** + +```rust +pub mod mapped_pinned; +``` + +- [ ] **Step 4: Delete the local copies from distributional_q_tests.rs and import from shared module** + +```rust +use crate::cuda_pipeline::mapped_pinned::{MappedF32Buffer, MappedI32Buffer}; +``` + +- [ ] **Step 5: Verify Test 0.F still runs end-to-end with bit-identical output** + +```bash +SQLX_OFFLINE=true cargo test -p ml --lib distributional_q_tests::test_0f -- --ignored --nocapture 2>&1 | tail -10 +``` + +Expected: σ_C51 = `[141.74062, 142.49536, 140.4468, 140.72096]`, argmax=2, P(active)=0.4330, panic with bias-hypothesis assertion. Bit-identical to commit `da632446c`'s output. + +- [ ] **Step 6: Commit** + +```bash +git add crates/ml/src/cuda_pipeline/mapped_pinned.rs \ + crates/ml/src/cuda_pipeline/mod.rs \ + crates/ml/src/trainers/dqn/distributional_q_tests.rs +git commit -m "refactor(cuda_pipeline): promote MappedF32/I32Buffer to shared module + +Moves MappedF32Buffer + MappedI32Buffer from distributional_q_tests.rs +local definitions to crates/ml/src/cuda_pipeline/mapped_pinned.rs so all +kernel test wrappers (Phase 0.F, upcoming MoE tests) share one +implementation. + +Adds write_from_slice helper for direct host_ptr write (no memcpy). + +Test 0.F bit-identical: sigma_C51, argmax, Thompson counts unchanged. + +Per feedback_no_htod_htoh_only_mapped_pinned.md: this is the only +allowed CPU<->GPU path. The shared module makes the mapped-pinned +pattern more discoverable and reduces duplication. + +Co-Authored-By: Claude Opus 4.7 (1M context) " +``` ### Task 2.1 — moe_mixture_forward kernel @@ -958,7 +1044,9 @@ impl GpuMoeHead { }) } - /// Test-only entry: upload host buffers, run kernel, download result. + /// Test-only entry: stage CPU-side data into mapped pinned memory, + /// run kernel, read GPU-written output via mapped pinned. **No HtoD, + /// no HtoH** — `feedback_no_htod_htoh_only_mapped_pinned.md`. pub fn test_mixture_forward( &self, expert_outputs: &[f32], @@ -967,9 +1055,16 @@ impl GpuMoeHead { k: usize, c: usize, ) -> Result, MLError> { - let expert_dev: CudaSlice = self.stream.memcpy_stod(expert_outputs)?; - let gate_dev: CudaSlice = self.stream.memcpy_stod(gate)?; - let mut h_s2_dev: CudaSlice = self.stream.alloc_zeros(b * c)?; + // Allocate mapped pinned buffers (CPU-and-GPU coherent, no copy). + // Pattern: cuMemHostAlloc(DEVICEMAP|PORTABLE) → cuMemHostGetDevicePointer_v2. + // Reference: MappedF32Buffer in crates/ml/src/trainers/dqn/distributional_q_tests.rs. + let expert_buf = unsafe { mapped_f32::MappedF32Buffer::new(k * b * c) }?; + let gate_buf = unsafe { mapped_f32::MappedF32Buffer::new(b * k) }?; + let h_s2_buf = unsafe { mapped_f32::MappedF32Buffer::new(b * c) }?; + + // Write CPU-side via host_ptr (no memcpy, just direct memory writes). + expert_buf.write_from_slice(expert_outputs); + gate_buf.write_from_slice(gate); let total = b * c; let block = 256_u32; @@ -978,9 +1073,9 @@ impl GpuMoeHead { unsafe { self.stream .launch_builder(&self.moe_mixture_forward) - .arg(&expert_dev) - .arg(&gate_dev) - .arg(&mut h_s2_dev) + .arg(&expert_buf.dev_ptr) + .arg(&gate_buf.dev_ptr) + .arg(&h_s2_buf.dev_ptr) .arg(&(b as i32)) .arg(&(k as i32)) .arg(&(c as i32)) @@ -991,15 +1086,25 @@ impl GpuMoeHead { }) .map_err(|e| MLError::ModelError(format!("moe_mixture_forward launch: {e}")))?; } + // Stream sync ensures __threadfence_system() side effects are visible + // to the host_ptr read below. self.stream .synchronize() .map_err(|e| MLError::ModelError(format!("synchronize: {e}")))?; - let result = self.stream.memcpy_dtov(&h_s2_dev) - .map_err(|e| MLError::ModelError(format!("download: {e}")))?; - Ok(result) + // Read GPU-written result via host_ptr (mapped pinned coherence). + Ok(h_s2_buf.read_all()) } } + +/// Promote the existing `MappedF32Buffer` from distributional_q_tests.rs +/// to a shared module so all kernel test wrappers share one implementation. +/// Lives at `crates/ml/src/cuda_pipeline/mapped_pinned.rs` and re-exports +/// `MappedF32Buffer`, `MappedI32Buffer`. The `distributional_q_tests` +/// module's local copies are deleted in favor of the shared module. +mod mapped_f32 { + pub use crate::cuda_pipeline::mapped_pinned::MappedF32Buffer; +} ``` - [ ] **Step 6: Add module to mod.rs** diff --git a/docs/superpowers/specs/2026-04-27-moe-regime-redesign-design.md b/docs/superpowers/specs/2026-04-27-moe-regime-redesign-design.md index 08753993c..440fb7b80 100644 --- a/docs/superpowers/specs/2026-04-27-moe-regime-redesign-design.md +++ b/docs/superpowers/specs/2026-04-27-moe-regime-redesign-design.md @@ -181,6 +181,10 @@ All graph-capturable; cuBLAS calls capture; existing `softmax_forward`/`softmax_ Forward graph extends with new nodes for: gate forward (2 SGEMMs + softmax), expert forward (8 pairs of SGEMMs + LeakyReLU), `moe_mixture_forward`, `moe_load_balance_loss`. Backward graph extends symmetrically. ISV producer launches in the cold-path cadence (one per training step). Graph recaptures cleanly on lr change / fold boundary per existing pattern. +### 6.4 Memory-transfer constraint: mapped pinned only + +Per `feedback_no_htod_htoh_only_mapped_pinned.md`: every CPU↔GPU communication path in the MoE implementation uses **mapped pinned memory** (`cuMemHostAlloc(DEVICEMAP | PORTABLE)` + `cuMemHostGetDevicePointer_v2`). No `cudaMemcpy` HtoD, no Vec-to-Vec defensive copies, including in test code. The reference template is `MappedF32Buffer` / `MappedI32Buffer` from `crates/ml/src/trainers/dqn/distributional_q_tests.rs` (used by Test 0.F); the implementation plan promotes these to a shared `crates/ml/src/cuda_pipeline/mapped_pinned.rs` module so all kernel test wrappers and any future CPU→GPU staging share one implementation. CPU is strictly read-only on the production training surface — kernel-arg scalars (passed through the launch ABI by value) are the only host-originated values that flow into kernels; everything else either originates GPU-side (replay sample, ISV slot, prior kernel output) or reaches the kernel via a mapped pinned `dev_ptr`. + ## 7. Cleanup cascade (atomic, single commit) Per `feedback_no_partial_refactor.md`, every consumer of the deleted infrastructure migrates in the same commit: