diff --git a/crates/ml-alpha/cuda/gpu_sample_and_gather.cu b/crates/ml-alpha/cuda/gpu_sample_and_gather.cu index 14f47240a..c897c7809 100644 --- a/crates/ml-alpha/cuda/gpu_sample_and_gather.cu +++ b/crates/ml-alpha/cuda/gpu_sample_and_gather.cu @@ -24,6 +24,8 @@ #define BOOK_LEVELS 10 #define REGIME_DIM 6 +#define N_HORIZONS 3 +#define N_LABEL_SOURCES 5 // Must match #[repr(C)] Mbp10RawInput in snap_features.rs (216 bytes). struct __align__(8) Mbp10Raw { @@ -80,11 +82,26 @@ extern "C" __global__ void gpu_sample_and_gather( long long* __restrict__ ts_ns_soa, // [B*K] long long* __restrict__ prev_ts_ns_soa, // [B*K] // Global-memory outputs for the sampled (file_offset, anchor, file_idx) - // per batch element. Read by gpu_gather_next, gpu_gather_frd_labels, - // and gpu_gather_bce_labels to reuse the same sampling decision. + // per batch element. Read by gpu_gather_next and gpu_gather_frd_labels + // to reuse the same sampling decision. int* __restrict__ out_file_offset, // [B] int* __restrict__ out_anchor, // [B] int* __restrict__ out_file_idx, // [B] + // Fused label gather — 5 label sources (BCE, prof_long, prof_short, + // size_long, size_short), each [N_HORIZONS × total_snaps]. Reads in + // the same pass as snapshot gather to avoid 5 separate random-access + // kernel launches that consumed 95.6% of GPU time. + int total_snaps, + const float* __restrict__ labels_src_0, // [N_HORIZONS × total_snaps] + const float* __restrict__ labels_src_1, + const float* __restrict__ labels_src_2, + const float* __restrict__ labels_src_3, + const float* __restrict__ labels_src_4, + float* __restrict__ labels_out_0, // [K, B, N_HORIZONS] + float* __restrict__ labels_out_1, + float* __restrict__ labels_out_2, + float* __restrict__ labels_out_3, + float* __restrict__ labels_out_4, int B ) { int b = blockIdx.x; @@ -148,9 +165,26 @@ extern "C" __global__ void gpu_sample_and_gather( tc_soa[n] = (int)snap.trade_count; ts_ns_soa[n] = (long long)snap.ts_ns; prev_ts_ns_soa[n] = (long long)snap.prev_ts_ns; + + // Fused label gather — same global_idx, zero extra random access. + // Output layout: [K, B, N_HORIZONS] row-major. + int lbl_base = (k * B + b) * N_HORIZONS; + const float* srcs[N_LABEL_SOURCES] = { + labels_src_0, labels_src_1, labels_src_2, labels_src_3, labels_src_4 + }; + float* outs[N_LABEL_SOURCES] = { + labels_out_0, labels_out_1, labels_out_2, labels_out_3, labels_out_4 + }; + #pragma unroll + for (int s = 0; s < N_LABEL_SOURCES; s++) { + #pragma unroll + for (int h = 0; h < N_HORIZONS; h++) { + outs[s][lbl_base + h] = srcs[s][h * total_snaps + global_idx]; + } + } } -// ── Label gather kernel ────────────────────────────────────────────── +// ── Label gather kernel (STANDALONE — kept for backward compat) ───── // // Runs AFTER gpu_sample_and_gather. Gathers per-horizon FRD labels at // the anchor position (rightmost K position = newest snapshot in the diff --git a/crates/ml-alpha/src/data/gpu_dataset.rs b/crates/ml-alpha/src/data/gpu_dataset.rs index 66d29fe62..0e34b67a8 100644 --- a/crates/ml-alpha/src/data/gpu_dataset.rs +++ b/crates/ml-alpha/src/data/gpu_dataset.rs @@ -188,6 +188,7 @@ impl GpuDataLoader { soa: &SoaBufferPtrs, seq_len: usize, batch_size: usize, + label_outs: [u64; 5], ) -> Result<()> { let mut args = RawArgs::new(); args.push_ptr(dataset.snapshots_d.raw_ptr()); @@ -210,6 +211,18 @@ impl GpuDataLoader { args.push_ptr(self.sample_file_offset_d.raw_ptr()); args.push_ptr(self.sample_anchor_d.raw_ptr()); args.push_ptr(self.sample_file_idx_d.raw_ptr()); + // Fused label gather arguments + args.push_i32(dataset.total_snapshots as i32); + args.push_ptr(dataset.labels_d.raw_ptr()); + args.push_ptr(dataset.outcome_prof_long_d.raw_ptr()); + args.push_ptr(dataset.outcome_prof_short_d.raw_ptr()); + args.push_ptr(dataset.outcome_size_long_d.raw_ptr()); + args.push_ptr(dataset.outcome_size_short_d.raw_ptr()); + args.push_ptr(label_outs[0]); + args.push_ptr(label_outs[1]); + args.push_ptr(label_outs[2]); + args.push_ptr(label_outs[3]); + args.push_ptr(label_outs[4]); args.push_i32(batch_size as i32); let mut ptrs = args.build_arg_ptrs(); unsafe { diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index c46fdb99a..a04b284f2 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -6849,8 +6849,15 @@ impl IntegratedTrainer { // 3. forward_encoder_from_device: runs encoder on SoA (= s_{t+1}). // 4. DtoD h_t -> h_tp1: saves h_{t+1}. let soa = self.perception.soa_buffer_ptrs(); + let label_stg = self.perception.label_staging_ptrs(); gpu_loader - .sample_and_gather(gpu_dataset, &soa, seq_len, b_size) + .sample_and_gather(gpu_dataset, &soa, seq_len, b_size, [ + label_stg.labels, + label_stg.aux_y_prof_long, + label_stg.aux_y_prof_short, + label_stg.aux_y_size_long, + label_stg.aux_y_size_short, + ]) .context("step_with_lobsim_gpu: sample_and_gather")?; gpu_loader .gather_next_window(gpu_dataset, &soa, seq_len, b_size) @@ -6895,47 +6902,10 @@ impl IntegratedTrainer { // After this call the encoder weights are updated with BCE+aux // gradient signal, so the subsequent RL forward on h_t sees the // jointly-supervised encoder. + // Label gathers fused into sample_and_gather above — 5 separate + // random-access launches eliminated (was 95.6% of GPU time per nsys). { use crate::heads::N_HORIZONS as PERC_N_HORIZONS; - let stg = self.perception.label_staging_ptrs(); - let total_snaps = gpu_dataset.total_snapshots; - - // BCE labels → stg_labels - gpu_loader - .gather_bce_labels( - gpu_dataset.labels_d.raw_ptr(), total_snaps, - PERC_N_HORIZONS, seq_len, b_size, stg.labels, - ) - .context("step_with_lobsim_gpu: gather_bce_labels")?; - // Aux prof-long → stg_aux_y_prof_long - gpu_loader - .gather_bce_labels( - gpu_dataset.outcome_prof_long_d.raw_ptr(), total_snaps, - PERC_N_HORIZONS, seq_len, b_size, stg.aux_y_prof_long, - ) - .context("step_with_lobsim_gpu: gather outcome_prof_long")?; - // Aux prof-short → stg_aux_y_prof_short - gpu_loader - .gather_bce_labels( - gpu_dataset.outcome_prof_short_d.raw_ptr(), total_snaps, - PERC_N_HORIZONS, seq_len, b_size, stg.aux_y_prof_short, - ) - .context("step_with_lobsim_gpu: gather outcome_prof_short")?; - // Aux size-long → stg_aux_y_size_long - gpu_loader - .gather_bce_labels( - gpu_dataset.outcome_size_long_d.raw_ptr(), total_snaps, - PERC_N_HORIZONS, seq_len, b_size, stg.aux_y_size_long, - ) - .context("step_with_lobsim_gpu: gather outcome_size_long")?; - // Aux size-short → stg_aux_y_size_short - gpu_loader - .gather_bce_labels( - gpu_dataset.outcome_size_short_d.raw_ptr(), total_snaps, - PERC_N_HORIZONS, seq_len, b_size, stg.aux_y_size_short, - ) - .context("step_with_lobsim_gpu: gather outcome_size_short")?; - // pos_fraction: gathered into a small device scratch, then // read back on the host to compute clamped pos_weight. The // gather writes to a mapped-pinned buffer so we can read it