Files

jgrusewski 4709ca8bc2 feat(dqn): enable Branching DQN with 45 factored actions (5×3×3)

Restore 45-action factored space via Branching DQN (Tavakoli 2018),
outputting 11 Q-values (5+3+3) instead of 45. This was reduced to 5
exposure-only actions during debugging and was never intended as permanent.

- Enable use_branching: true by default in DQNConfig and DQNHyperparameters
- Add branching paths to select_action_with_confidence and select_action_inference
- Update agent.rs select_action_factored for branching-aware selection
- Expand CountBonus to per-branch tracking with bonuses_branched()
- Add order_type + urgency distribution tracking in monitoring
- Add DQN_ORDER_ACTIONS=3, DQN_URGENCY_ACTIONS=3, DQN_TOTAL_ACTIONS=45 to CUDA header
- Fix 7 pre-existing clippy doc_markdown errors in regime_conditional.rs
- Fix pre-existing cognitive_complexity in replay_buffer_type.rs (extract helpers)
- Fix flaky GPU test OOM under parallel execution (CPU fallback + test VRAM safety)
- Delete unused flash_attention submodules (block_sparse, causal_masking, etc.)
- Add GPU hot-path guard scripts and ensemble/hyperopt adapter improvements

Tests: ml-dqn 416/0, ml 905/0, clippy 0 errors on both crates

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

2026-03-11 22:00:13 +01:00

8.9 KiB

Raw Blame History

H100 CUDA Optimization Hive Plan

Overview

Hive-mind swarm execution plan for 14 H100-targeted CUDA optimizations across 5 phases. Target: 8-16x faster experience collection, 3-5x end-to-end training speedup on H100 SXM5.

H100 SXM5 Reference Specs

Spec	Value	Current Codebase
HBM3	80 GB @ 3.35 TB/s	VRAM scaling [2048,1024,512]
BF16 Tensor Core	989 TFLOPS	BF16 default via training_dtype()
FP8 E4M3	1,979 TFLOPS	NOT USED
SMs	132	optimal_n_episodes() aware
L2 Cache	50 MB	NO PERSISTENCE HINTS
Shared Mem/SM	228 KB	64 KB tiles (28%)
Max Threads/SM	2,048	64 used (3% occupancy)

Swarm Topology

                        QUEEN COORDINATOR
                              |
          +-------------------+-------------------+
          |                   |                   |
      WAVE 1 (4)          WAVE 2 (2)         VALIDATOR
     [parallel]          [sequential]        [continuous]
          |                   |
      WAVE 3 (2)          WAVE 4 (2)
     [parallel]          [parallel]

Topology: Hierarchical with wave-gated phases. Max concurrent agents: 4 (Wave 1), scales to 2 in later waves. Isolation: Each agent operates in a dedicated git worktree.

Wave Execution Plan

Wave 0: Measurement Baseline

Agent: profiling-scout

Branch: feature/h100-nvtx
Task: Add NVTX range markers to all CUDA paths. Add cudaMallocAsync pool wrapper.
Files: cuda_pipeline/*.rs, ml-dqn/src/gpu_replay_buffer.rs
Deliverable: NVTX-instrumented build, nsys profiling guide
Acceptance: cargo test passes, NVTX ranges visible in nsys profile

Wave 1: Quick Wins (Parallel, No Dependencies)

Agent: l2-cache-worker

Branch: feature/h100-l2-pinning
Task: cudaAccessPolicyWindow for model weights. H100 detection via gpu_name.
Files: gpu_weights.rs, gpu_experience_collector.rs
LOC estimate: ~20 lines
Acceptance: cargo test passes, no regression

Agent: shmem-worker

Branch: feature/h100-shmem-228k
Task: SHMEM_TILE_ROWS=128 via NVRTC define. cudaFuncSetAttribute(228KB).
Files: common_device_functions.cuh, gpu_experience_collector.rs
LOC estimate: ~5 lines + NVRTC injection
Acceptance: cargo test passes, tile loop count reduced

Agent: double-buffer-worker

Branch: feature/h100-async-dbuf
Task: Wire CudaStreamPool into DoubleBufferedLoader for async staging.
Files: double_buffer.rs, cuda_streams.rs
LOC estimate: ~30 lines
Acceptance: cargo test passes, staging overlaps training

Wave 2: Core Kernel Rewrite (Sequential, Foundational)

Agent: kernel-architect (THE BIG ONE)

Branch: feature/h100-warp-cooperative
Task: Port q_forward_dueling_dist's warp-distributed approach to standard q_forward_dueling. Distribute scratch arrays across 32 lanes (3KB/thread -> 96 bytes/thread). Use warp_reduce_sum_all for dot products, __shfl_sync for broadcasts. Add golden-reference comparison test (sequential vs warp, max delta < 1e-4).
Files: dqn_experience_kernel.cu, common_device_functions.cuh
Risk: HIGH -- numerical divergence from changed FP addition ordering
Mitigation: Golden-reference Q-value comparison on first 1000 episodes
Acceptance: Golden-reference test passes, cargo test passes, occupancy > 8 warps/SM

Agent: variant-porter (DEPENDS ON kernel-architect)

Branch: feature/h100-warp-all-variants
Task: Extend warp-cooperative to NoisyNet + Branching DQN forward paths.
Files: dqn_experience_kernel.cu (noisy_matvec, branching sections)
Acceptance: Golden-reference tests for all variants, cargo test passes

Wave 3: Hopper-Specific Features (Parallel, After Wave 2)

Agent: tma-worker

Branch: feature/h100-tma
Task: Replace cooperative_load_tile() with TMA cp.async.bulk.tensor. Add sm_90 architecture guard, fallback to float4 loads on non-Hopper.
Files: common_device_functions.cuh, gpu_experience_collector.rs
Risk: MEDIUM -- Hopper-exclusive PTX, must degrade gracefully
Acceptance: cargo test passes on both H100 and non-H100 paths

Agent: fp8-researcher

Branch: feature/h100-fp8-inference
Task: FP8 E4M3 weight quantization for experience kernel only. __nv_fp8_e4m3 types in matvec kernels, per-tensor scale factors. Training stays BF16. A/B test infrastructure for Sharpe comparison.
Files: dqn_experience_kernel.cu, gpu_weights.rs, mixed_precision.rs
Risk: HIGH -- quantization may degrade policy quality
Mitigation: A/B Sharpe ratio comparison, reject if > 5% degradation
Acceptance: A/B Sharpe delta < 5%, cargo test passes

Wave 4: Orchestration (Parallel, After Waves 2-3)

Agent: graph-worker

Branch: feature/h100-cuda-graphs
Task: CUDA Graph capture for backtest step loop and training step. Pre-capture graph pool for batch size tiers (128, 256, 512, 1024). Wire existing use_cuda_graphs config field.
Files: gpu_backtest_evaluator.rs, trainer.rs
Risk: MEDIUM -- fixed tensor shapes assumption vs adaptive allocator
Mitigation: Graph pool per batch size tier
Acceptance: cargo test passes, nsys shows single graph replay

Agent: pipeline-worker

Branch: feature/h100-multi-stream
Task: Pipeline training stages across 2-3 CUDA streams with cudaEvent sync. Evaluate persistent kernel pattern as alternative for experience collection.
Files: trainer.rs, cuda_streams.rs, gpu_experience_collector.rs
Acceptance: cargo test passes, SM utilization > 60% in nsys

Continuous: Validator Agent

No dedicated worktree -- reads from merged branches
After each wave merge: full test suite + clippy + gpu-hotpath-guard
Q-value distribution comparison before/after each merge
Gate: 905 DQN tests + 2758 total tests + 0 clippy warnings + 0 guard violations

Memory Architecture (Shared State)

hive/h100/
  phase           -> current active phase (0-4)
  wave1/
    l2/status     -> pending | running | done | failed
    shmem/status  -> pending | running | done | failed
    dbuf/status   -> pending | running | done | failed
  wave2/
    kernel-architect/status
    variant-porter/status
  wave3/
    tma/status
    fp8/status
  wave4/
    graphs/status
    pipeline/status
  validator/
    last-run      -> timestamp + pass/fail + test count
  golden-ref/
    q-values      -> baseline Q-value distribution hash
  metrics/
    occupancy     -> measured warps/SM after each wave
    throughput    -> measured exp/s after each wave

Queen Decision Gates

Gate	Condition	Action
0 -> 1	NVTX in place, nsys captured	Spawn Wave 1 agents (4 parallel)
1 -> 2	All quick wins merged + validated	Spawn kernel-architect
2 -> 3	Warp-cooperative merged, golden-ref delta < 1e-4	Spawn Wave 3 agents (2 parallel)
3 -> 4	Hopper features merged with fallback	Spawn Wave 4 agents (2 parallel)
4 -> Done	Full pipeline validated, benchmark improvement	Complete

Merge Protocol

Worker completes -> sets memory status to "ready-for-review"
Validator runs tests on worktree branch -> "validated" or "failed"
Queen merges validated worktrees into main (sequential, avoids conflicts)
Queen advances phase gate when all wave tasks merged + validated

Conflict Resolution

Wave	Conflict Risk	Resolution
Wave 1	NONE -- tasks touch different files	Parallel merge
Wave 2	HIGH -- both touch dqn_experience_kernel.cu	Sequential (architect first)
Wave 3	LOW -- TMA touches common_device_functions.cuh (shared)	Merge after Wave 2
Wave 4	MEDIUM -- both touch trainer.rs	Sequential merge, queen resolves

Risk Register

Risk	Severity	Mitigation
Warp-cooperative numerical divergence	HIGH	Golden-reference Q-value test, max delta < 1e-4
FP8 policy quality degradation	HIGH	A/B Sharpe ratio test, reject if > 5% delta
CUDA Graphs + dynamic batch sizes	MEDIUM	Pre-captured graph pool per batch tier
Hopper features on non-H100	MEDIUM	#ifdef CUDA_ARCH >= 900 + runtime detection
Test regression	LOW	Validator agent gates every wave merge
HBM fragmentation in long runs	LOW	cudaMallocAsync pool (Phase 0)

Expected Impact

Phase	Optimization	Metric
Wave 1	L2 cache pinning	3.5x less HBM read pressure
Wave 1	Larger shmem tiles	Eliminated tile loops for <=128-dim layers
Wave 1	Async double buffer	~0ms fold transition stall
Wave 2	Warp-cooperative forward	4-8x experience throughput
Wave 3	FP8 E4M3 inference	2x tensor core TFLOPS
Wave 3	TMA tile loads	Hidden memory load latency
Wave 4	CUDA Graphs	5-10x less launch overhead
Wave 4	Multi-stream pipeline	30-40% less idle SMs
Combined	All phases	8-16x experience collection, 3-5x training

Deferred

Thread Block Clusters: High complexity, uncertain ROI for prefix-sum. Re-evaluate after Wave 4 benchmarks.

8.9 KiB Raw Blame History