feat(rl): CUDA Graph B capture — post-fill reward/EMA/controller pipeline (~20 kernels)

Three-state machine (warmup → capture → replay) wraps the post-fill
section: extract_realized_pnl_delta, unit_state_update, trade_context,
multires, step_counter, 3× EMA, abs_copy, reward_shaping, raw_rewards
snapshot, apply_reward_scale, reward_clamp, atom_support_update,
compute_advantage_return, var_over_abs_mean.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-25 23:05:59 +02:00
parent dc7f4c6659
commit 3af2b40100

View File

@@ -425,6 +425,7 @@ pub struct IntegratedTrainer {
// barrier that can't be captured).
prefill_graph: Option<CudaGraph>,
postfill_graph: Option<CudaGraph>,
reward_graph: Option<CudaGraph>,
graph_warmup_done: bool,
// ── Kernel handles for grad combine + cross-batch reduce ──────────
@@ -1902,6 +1903,7 @@ impl IntegratedTrainer {
stream,
prefill_graph: None,
postfill_graph: None,
reward_graph: None,
graph_warmup_done: false,
_grad_h_module: grad_h_module,
grad_h_accumulate_fn,
@@ -5042,6 +5044,24 @@ impl IntegratedTrainer {
.step_fill_from_market_targets(last_snap.ts_ns)
.context("step_with_lobsim: lobsim.step_fill_from_market_targets")?;
// ── Graph B: post-fill reward/EMA/controller pipeline ─────────
// ~20 kernels from extract_realized_pnl_delta through
// var_over_abs_mean. Same three-state machine as Graph A/A2 —
// captures after A2 is done, replays on subsequent steps.
if self.reward_graph.is_some() {
self.reward_graph
.as_ref()
.unwrap()
.launch()
.context("reward graph launch")?;
} else {
let capturing_reward = self.postfill_graph.is_some();
if capturing_reward {
self.stream
.begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)
.map_err(|e| anyhow::anyhow!("reward begin_capture: {e}"))?;
}
// Extract per-batch (reward, done) deltas vs the trainer's
// prev_realized_pnl_d snapshot. Updates prev_* in place for
// the next step's delta.
@@ -5493,6 +5513,18 @@ impl IntegratedTrainer {
self.launch_var_over_abs_mean(&self.advantages_d.clone(), b_size)
.context("launch_var_over_abs_mean(advantages_d)")?;
if capturing_reward {
let graph = self
.stream
.end_capture(
CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
)
.context("reward end_capture")?
.ok_or_else(|| anyhow::anyhow!("reward end_capture returned None"))?;
self.reward_graph = Some(graph);
}
} // end else (reward warmup / capture)
// ── Step 7a (R7d): PER push + sample + gather. ────────────────
// Push CURRENT step's b_size transitions (one per batch) to
// ReplayBuffer; sample b_size indices (priority^α from