gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(per-horizon-cfc): wire Phase 2 fused CfC forward dispatch in trainer

Browse Source 
Per spec §2.1 + §5.4 point 1 and Task 10 of the plan.

Adds binding `cfc_step_per_branch_fwd_gpu` in cfc/step.rs that wraps
the fused per-branch kernel with the spec-mandated launch config:
grid=(B, N_HORIZONS=5, 1), block=(MAX_BUCKET_DIM=28, 1, 1) uniform
predicate, cooperative staging of x_local + h_old_local (2 × HIDDEN_DIM
floats). Also adds `cfc_step_per_branch_bwd_gpu` binding for Task 11.

CfcTrunk loads two new cubins (cfc_step_per_branch + heads_block_diagonal)
and caches three function handles. `heads_w_skip_compact_d` (from Task 9)
remains allocated as Phase-2-prepared metadata; full heads consumption
deferred to a follow-up task (Phase 2 heads dispatch needs GRN
integration, which is out of Task 10 scope per
`feedback_no_partial_refactor`).

In `dispatch_train_step`, the CfC dispatch branches on TrainingPhase:
- Phase1Warmup: existing `cfc_step_batched` (unchanged)
- Phase2Routed: `cfc_step_per_branch_fwd` consuming bucket-routed
  `tau_all_d` + `bucket_channel_offset_d` + `bucket_dim_k_d`

GRN heads dispatch stays unchanged for both phases per the Task 10
scope adjustment (Option B in the task brief). The `heads_w_skip` storage
remains 640 floats; the compact 128-float `heads_w_skip_compact_d` is
metadata-ready but not yet consumed.

Workspace cargo check clean (ml-alpha lib + trainer compile; only
pre-existing cupti/sp15/gpu_per_integration test errors remain).
ml-alpha lib tests: 33 passed.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-21 17:22:21 +02:00
parent 2f69bb1fc0
commit d16cab18bb
3 changed files with 212 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

108
crates/ml-alpha/src/cfc/step.rs
View File

@@ -89,6 +89,114 @@ pub fn cfc_step_backward_gpu(
))
}
/// Fused per-branch CfC forward step (Phase 2 dispatch).
///
/// Per spec §5.4 point 1 (docs/superpowers/specs/2026-05-21-per-horizon-cfc-inference-design.md):
/// single launch covers all 5 branches × n_batch with grid=(B, N_HORIZONS, 1)
/// and block=(MAX_BUCKET_DIM=28, 1, 1). Uniform-predicate handles uneven
/// bucket sizes [25, 25, 25, 25, 28] without warp divergence.
///
/// Buffer contract (matches kernel signature in
/// `crates/ml-alpha/cuda/cfc_step_per_branch.cu`):
/// - `w_in_d`, `w_rec_d`: shared `[HIDDEN_DIM × HIDDEN_DIM]` body weights.
/// - `b_d`: `[HIDDEN_DIM]` bias.
/// - `tau_all_d`: `[HIDDEN_DIM]` bucket-grouped τ (populated by transition).
/// - `x_d`, `h_old_d`: `[B × HIDDEN_DIM]` per-batch inputs.
/// - `bucket_channel_offset_d`: `[N_HORIZONS + 1]` cumulative bucket starts.
/// - `bucket_dim_k_d`: `[N_HORIZONS]` per-bucket channel counts.
/// - `h_new_d`: `[B × HIDDEN_DIM]` output (overwrite).
///
/// The cooperative-staging shared-memory layout (`2 × HIDDEN_DIM` floats)
/// stages `x` and `h_old` rows once per block per
/// `pearl_cooperative_staging_eliminates_redundant_reads`.
pub fn cfc_step_per_branch_fwd_gpu(
stream: &Arc<CudaStream>,
func: &cudarc::driver::CudaFunction,
w_in_d: &CudaSlice<f32>,
w_rec_d: &CudaSlice<f32>,
b_d: &CudaSlice<f32>,
tau_all_d: &CudaSlice<f32>,
x_d: &CudaSlice<f32>,
h_old_d: &CudaSlice<f32>,
dt_s: f32,
b_sz: i32,
bucket_channel_offset_d: &CudaSlice<u32>,
bucket_dim_k_d: &CudaSlice<u32>,
h_new_d: &mut CudaSlice<f32>,
) -> Result<()> {
const N_HORIZONS: u32 = 5;
const MAX_BUCKET_DIM: u32 = 28;
const HIDDEN_DIM: u32 = 128;
let cfg = LaunchConfig {
grid_dim: (b_sz as u32, N_HORIZONS, 1),
block_dim: (MAX_BUCKET_DIM, 1, 1),
// 2 × HIDDEN_DIM floats: x_local + h_old_local cooperative staging.
shared_mem_bytes: HIDDEN_DIM * 4 * 2,
};
let mut launch = stream.launch_builder(func);
launch
.arg(w_in_d).arg(w_rec_d).arg(b_d).arg(tau_all_d)
.arg(x_d).arg(h_old_d).arg(&dt_s).arg(&b_sz)
.arg(bucket_channel_offset_d).arg(bucket_dim_k_d)
.arg(h_new_d);
unsafe { launch.launch(cfg).context("cfc_step_per_branch_fwd launch")?; }
Ok(())
}
/// Fused per-branch CfC backward step (Phase 2 dispatch).
///
/// Mirror of `cfc_step_per_branch_fwd_gpu`. Per-batch grad slices are
/// written by this kernel; cross-batch reduction is the caller's
/// responsibility (matches existing `reduce_axis0` pattern in
/// `cfc_step_backward_batched`). Each thread (batch, branch, c) is the
/// sole writer of its slice — no atomicAdd per
/// `feedback_no_atomicadd`.
///
/// Note: this binding is staged for Task 11's backward dispatch wiring.
/// Task 10 only calls the forward path; Task 11 will land the backward
/// caller. The kernel itself is already validated by the GPU oracle
/// test in `crates/ml-alpha/tests/cfc_step_per_branch.rs`.
#[allow(clippy::too_many_arguments)]
pub fn cfc_step_per_branch_bwd_gpu(
stream: &Arc<CudaStream>,
func: &cudarc::driver::CudaFunction,
w_in_d: &CudaSlice<f32>,
w_rec_d: &CudaSlice<f32>,
b_d: &CudaSlice<f32>,
tau_all_d: &CudaSlice<f32>,
x_d: &CudaSlice<f32>,
h_old_d: &CudaSlice<f32>,
grad_h_new_d: &CudaSlice<f32>,
dt_s: f32,
b_sz: i32,
bucket_channel_offset_d: &CudaSlice<u32>,
bucket_dim_k_d: &CudaSlice<u32>,
grad_w_in_d: &mut CudaSlice<f32>,
grad_w_rec_d: &mut CudaSlice<f32>,
grad_b_d: &mut CudaSlice<f32>,
grad_tau_all_d: &mut CudaSlice<f32>,
grad_h_old_d: &mut CudaSlice<f32>,
) -> Result<()> {
const N_HORIZONS: u32 = 5;
const MAX_BUCKET_DIM: u32 = 28;
const HIDDEN_DIM: u32 = 128;
let cfg = LaunchConfig {
grid_dim: (b_sz as u32, N_HORIZONS, 1),
block_dim: (MAX_BUCKET_DIM, 1, 1),
shared_mem_bytes: HIDDEN_DIM * 4 * 2,
};
let mut launch = stream.launch_builder(func);
launch
.arg(w_in_d).arg(w_rec_d).arg(b_d).arg(tau_all_d)
.arg(x_d).arg(h_old_d).arg(grad_h_new_d)
.arg(&dt_s).arg(&b_sz)
.arg(bucket_channel_offset_d).arg(bucket_dim_k_d)
.arg(grad_w_in_d).arg(grad_w_rec_d).arg(grad_b_d)
.arg(grad_tau_all_d).arg(grad_h_old_d);
unsafe { launch.launch(cfg).context("cfc_step_per_branch_bwd launch")?; }
Ok(())
}
pub fn cfc_step_gpu(
dev: &MlDevice,
w: &CfcWeights,

43
crates/ml-alpha/src/cfc/trunk.rs
View File

@@ -80,6 +80,11 @@ const HEADS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/multi_horiz
const LAYER_NORM_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/layer_norm.cubin"));
const VARIABLE_SELECTION_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/variable_selection.cubin"));
const ATTENTION_POOL_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/attention_pool.cubin"));
// Per-horizon CfC Phase 2 dispatch (spec §5.4 points 1–2). Modules
// stay alive on the trunk so the cached `CudaFunction` handles below
// remain valid through the trainer's lifetime.
const CFC_STEP_PER_BRANCH_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/cfc_step_per_branch.cubin"));
const HEADS_BLOCK_DIAGONAL_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/heads_block_diagonal_fwd.cubin"));
#[derive(Clone, Debug)]
pub struct CfcConfig {
@@ -126,6 +131,10 @@ pub struct CfcTrunk {
_ln_module: Arc<CudaModule>,
_vsn_module: Arc<CudaModule>,
_attn_module: Arc<CudaModule>,
// Per-horizon CfC Phase 2 cubin modules. Kept alive alongside the
// cached function handles below.
_cfc_step_per_branch_module: Arc<CudaModule>,
_heads_block_diagonal_module: Arc<CudaModule>,
// Forward kernel handles — read by PerceptionTrainer's evaluate_batched.
pub vsn_fwd_fn: CudaFunction,
pub ln_fwd_fn: CudaFunction,
@@ -134,6 +143,19 @@ pub struct CfcTrunk {
pub step_batched_fn: CudaFunction,
pub heads_grn_fwd_fn: CudaFunction,
pub transpose_3d_fn: CudaFunction,
// Per-horizon CfC Phase 2 dispatch handles (spec §5.4 points 1–2).
// `cfc_step_per_branch_fwd_fn` is read by `dispatch_train_step`'s
// Phase 2 branch (Task 10) and the eventual `forward_step_into`
// Phase 2 path (Task 13). The `_bwd_fn` handle is staged here so
// Task 11's backward wiring lands without re-touching the trunk
// constructor. `heads_block_diagonal_fwd_fn` is loaded for the
// future compact-w_skip heads dispatch; per Task 10's scope (see
// perception.rs Phase 2 dispatch comment), it is NOT consumed yet
// — heads continue through the existing GRN kernel until a later
// task wires the compact-storage path through the GRN's skip term.
pub cfc_step_per_branch_fwd_fn: CudaFunction,
pub cfc_step_per_branch_bwd_fn: CudaFunction,
pub heads_block_diagonal_fwd_fn: CudaFunction,
// CfC weights (cfc_n_in = HIDDEN_DIM). Public so PerceptionTrainer can
// mutate during init + AdamW.
@@ -192,6 +214,13 @@ impl CfcTrunk {
let ln_module = ctx.load_cubin(LAYER_NORM_CUBIN.to_vec()).context("load ln cubin")?;
let vsn_module = ctx.load_cubin(VARIABLE_SELECTION_CUBIN.to_vec()).context("load vsn cubin")?;
let attn_module = ctx.load_cubin(ATTENTION_POOL_CUBIN.to_vec()).context("load attn cubin")?;
// Per-horizon CfC Phase 2 cubin modules.
let cfc_step_per_branch_module = ctx
.load_cubin(CFC_STEP_PER_BRANCH_CUBIN.to_vec())
.context("load cfc_step_per_branch cubin")?;
let heads_block_diagonal_module = ctx
.load_cubin(HEADS_BLOCK_DIAGONAL_CUBIN.to_vec())
.context("load heads_block_diagonal cubin")?;
let vsn_fwd_fn = vsn_module.load_function("variable_selection_fwd").context("vsn_fwd_fn")?;
let ln_fwd_fn = ln_module.load_function("layer_norm_fwd").context("ln_fwd_fn")?;
let attn_fwd_fn = attn_module.load_function("attention_pool_fwd").context("attn_fwd_fn")?;
@@ -199,6 +228,15 @@ impl CfcTrunk {
let step_batched_fn = step_module.load_function("cfc_step_batched").context("step_batched_fn")?;
let heads_grn_fwd_fn = heads_module.load_function("multi_horizon_heads_grn_fwd_batched").context("heads_grn_fwd_fn")?;
let transpose_3d_fn = step_module.load_function("transpose_3d_swap_01").context("transpose_3d_fn")?;
let cfc_step_per_branch_fwd_fn = cfc_step_per_branch_module
.load_function("cfc_step_per_branch_fwd")
.context("cfc_step_per_branch_fwd_fn")?;
let cfc_step_per_branch_bwd_fn = cfc_step_per_branch_module
.load_function("cfc_step_per_branch_bwd")
.context("cfc_step_per_branch_bwd_fn")?;
let heads_block_diagonal_fwd_fn = heads_block_diagonal_module
.load_function("heads_block_diagonal_fwd")
.context("heads_block_diagonal_fwd_fn")?;
// CfC weight init via a seeded ChaCha8Rng — the same chain feeds
// every weight tensor so cfg.seed fully determines the initial
@@ -317,6 +355,8 @@ impl CfcTrunk {
_ln_module: ln_module,
_vsn_module: vsn_module,
_attn_module: attn_module,
_cfc_step_per_branch_module: cfc_step_per_branch_module,
_heads_block_diagonal_module: heads_block_diagonal_module,
vsn_fwd_fn,
ln_fwd_fn,
attn_fwd_fn,
@@ -324,6 +364,9 @@ impl CfcTrunk {
step_batched_fn,
heads_grn_fwd_fn,
transpose_3d_fn,
cfc_step_per_branch_fwd_fn,
cfc_step_per_branch_bwd_fn,
heads_block_diagonal_fwd_fn,
})
}

69
crates/ml-alpha/src/trainer/perception.rs
View File

@@ -2412,19 +2412,72 @@ impl PerceptionTrainer {
let main_k_ptr = main_base + (k * kb_nh_bytes) as u64;
let logit_k_ptr = logit_base + (k * kb_nh_bytes) as u64;
// cfc_step_batched(W, x[B*n_in], h_old[B*n_hid], …, → h_new[B*n_hid]).
// CfC forward dispatch, two-phase per spec §2.3 + Task 10:
//
// Phase 1 (warmup): existing single shared `cfc_step_batched`
// kernel — block-per-batch over all HIDDEN_DIM channels.
//
// Phase 2 (routed): fused per-branch kernel
// `cfc_step_per_branch_fwd` — grid=(B, N_HORIZONS=5, 1),
// block=(MAX_BUCKET_DIM=28, 1, 1) uniform predicate, with
// `bucket_channel_offset_d` / `bucket_dim_k_d` routing each
// block to its quintile of HIDDEN_DIM. Spec §5.4 point 1.
//
// The branch on `self.phase` lives OUTSIDE the captured graph's
// hot kernel scope: the captured train graph is recaptured at
// the Phase 1→2 transition (cf. `pearl_no_host_branches_in_captured_graph`).
// Once Phase 2 has been entered, every subsequent capture replays
// the routed path; the if-branch here is evaluated by host code
// building the graph, not inside the device-side kernel stream.
unsafe {
let mut launch = self.stream.launch_builder(&self.trunk.step_batched_fn);
launch
.arg(&self.trunk.w_in_d).arg(&self.trunk.w_rec_d).arg(&self.trunk.b_d).arg(&self.trunk.tau_all_d)
.arg(&x_k_ptr).arg(&h_old_k_ptr)
.arg(&dt_s).arg(&n_in_i).arg(&n_hid_i).arg(&n_batch_i)
.arg(&h_new_k_ptr);
launch.launch(cfg_cfc).context("cfc fwd k batched")?;
match self.phase {
TrainingPhase::Phase1Warmup => {
let mut launch = self.stream.launch_builder(&self.trunk.step_batched_fn);
launch
.arg(&self.trunk.w_in_d).arg(&self.trunk.w_rec_d).arg(&self.trunk.b_d).arg(&self.trunk.tau_all_d)
.arg(&x_k_ptr).arg(&h_old_k_ptr)
.arg(&dt_s).arg(&n_in_i).arg(&n_hid_i).arg(&n_batch_i)
.arg(&h_new_k_ptr);
launch.launch(cfg_cfc).context("cfc fwd k batched")?;
}
TrainingPhase::Phase2Routed => {
// grid=(B, N_HORIZONS, 1), block=(MAX_BUCKET_DIM, 1, 1).
// Shared mem = 2 × HIDDEN_DIM floats (x_local + h_old_local
// cooperative staging per pearl_cooperative_staging_eliminates_redundant_reads).
const N_HORIZONS_U32: u32 = N_HORIZONS as u32;
const MAX_BUCKET_DIM: u32 = 28;
let cfg_cfc_per_branch = LaunchConfig {
grid_dim: (b_sz as u32, N_HORIZONS_U32, 1),
block_dim: (MAX_BUCKET_DIM, 1, 1),
shared_mem_bytes: (2 * HIDDEN_DIM * std::mem::size_of::<f32>()) as u32,
};
let mut launch = self.stream.launch_builder(&self.trunk.cfc_step_per_branch_fwd_fn);
launch
.arg(&self.trunk.w_in_d).arg(&self.trunk.w_rec_d).arg(&self.trunk.b_d).arg(&self.trunk.tau_all_d)
.arg(&x_k_ptr).arg(&h_old_k_ptr)
.arg(&dt_s).arg(&n_batch_i)
.arg(&self.trunk.bucket_channel_offset_d).arg(&self.trunk.bucket_dim_k_d)
.arg(&h_new_k_ptr);
launch.launch(cfg_cfc_per_branch).context("cfc_step_per_branch_fwd k")?;
}
}
}
// multi_horizon_heads_grn_fwd_batched(10 params + h[B*128] + n_batch
// → probs[B*5] + z1/a1/z2 [B,5,HEAD_MID] + gate_logit/main/logit [B,5]).
//
// Task 10 scope note (Option B): the heads dispatch stays at the
// existing full-`heads_w_skip_d` GRN kernel for BOTH phases. The
// Phase 2 compact-w_skip path (`heads_block_diagonal_fwd`) and
// its 128-float `heads_w_skip_compact_d` storage are populated
// by the transition kernels (Task 9) but NOT yet consumed here:
// the GRN kernel computes the full per-horizon logit including
// skip + gate + main sub-components, and routing only the skip
// term through the compact path requires integration with
// GRN's main/gate computations (out of Task 10 scope per
// `feedback_no_partial_refactor`). A follow-up task wires the
// compact heads dispatch through GRN; the trunk handle
// `heads_block_diagonal_fwd_fn` is already loaded for that.
unsafe {
let mut launch = self.stream.launch_builder(&self.trunk.heads_grn_fwd_fn);
launch