entrenar/finetune/instruct_pipeline/
cuda_forward.rs

1#[cfg(feature = "cuda")]
2use super::{CudaBlockScratch, InstructGpuTrainingState, InstructPipeline, Transformer};
3
4#[cfg(feature = "cuda")]
5use crate::autograd::cuda_training::CudaTrainer;
6#[cfg(feature = "cuda")]
7use crate::transformer::CudaBlock;
8#[cfg(feature = "cuda")]
9use trueno_gpu::driver::{CaptureMode, GpuBuffer};
10
11#[cfg(feature = "cuda")]
12impl InstructPipeline {
13    /// GPU-accelerated forward pass saving layer inputs for backward.
14    #[allow(unsafe_code)]
15    pub(super) fn forward_cuda_training(
16        model: &Transformer,
17        token_ids: &[u32],
18        trainer: &CudaTrainer,
19        cuda_blocks: &mut [CudaBlock],
20        training_state: &mut InstructGpuTrainingState,
21        shared_scratch: &mut Option<CudaBlockScratch>,
22    ) -> Option<()> {
23        let seq_len = token_ids.len();
24        let hidden_size = model.config.hidden_size;
25        let max_seq_len = shared_scratch
26            .as_ref()
27            .map_or(model.config.max_position_embeddings.min(512), |s| s.max_seq_len(hidden_size));
28        let seq_len = if seq_len > max_seq_len { max_seq_len } else { seq_len };
29        if seq_len == 0 {
30            return None;
31        }
32
33        // Embed on CPU, upload to GPU
34        let hidden = model.embed_tokens.forward(token_ids);
35        let hidden_data = hidden.data();
36        let hidden_slice = hidden_data.as_slice().expect("contiguous hidden");
37
38        if crate::transformer::cuda_block::nan_scan_enabled() {
39            let bad = hidden_slice.iter().filter(|v| !v.is_finite()).count();
40            if bad > 0 {
41                eprintln!("[NAN-SCAN] embed (CPU): {bad}/{} non-finite", hidden_slice.len());
42            }
43        }
44
45        // PMAT-420 / entrenar#316: Use seq_len-sized fresh buffers (like inference forward).
46        training_state.fwd_scratch_a = trainer
47            .upload(hidden_slice)
48            .map_err(|e| eprintln!("[CUDA] embed upload failed: {e}"))
49            .ok()?;
50        training_state.fwd_scratch_b = trainer
51            .zeros(seq_len * hidden_size)
52            .map_err(|e| eprintln!("[CUDA] scratch_b alloc failed: {e}"))
53            .ok()?;
54
55        let scratch_a_ptr: *mut GpuBuffer<f32> =
56            std::ptr::from_mut(&mut training_state.fwd_scratch_a);
57        let scratch_b_ptr: *mut GpuBuffer<f32> =
58            std::ptr::from_mut(&mut training_state.fwd_scratch_b);
59        let mut input_is_a = true;
60
61        let stream = trainer.stream();
62        // entrenar#318: GPU-side scratch + training state zeroing (PMAT-453 NaN cascade fix).
63        if let Some(ref mut scratch) = shared_scratch.as_mut() {
64            scratch.zero_forward_buffers(stream);
65        }
66        for b in [
67            &mut training_state.grad_buf_a,
68            &mut training_state.grad_buf_b,
69            &mut training_state.grad_hidden_buf,
70            &mut training_state.output_scratch,
71            &mut training_state.logits_buf,
72        ] {
73            b.zero_async(stream).ok();
74        }
75
76        // PMAT-464: CUDA graph capture/replay (CUDA_GRAPH=1)
77        static USE_CUDA_GRAPH: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
78        let use_graph =
79            *USE_CUDA_GRAPH.get_or_init(|| std::env::var("CUDA_GRAPH").as_deref() == Ok("1"));
80
81        for (i, block_) in cuda_blocks.iter().enumerate() {
82            let _ = block_;
83            let expected_len = seq_len * hidden_size;
84            if training_state.layer_inputs[i].len() != expected_len {
85                training_state.layer_inputs[i] = trainer
86                    .zeros(expected_len)
87                    .map_err(|e| eprintln!("[CUDA] layer_input prealloc L{i}: {e}"))
88                    .ok()?;
89            }
90        }
91
92        if use_graph
93            && training_state.graph_cached_seq_len == seq_len
94            && training_state.forward_graph_exec.is_some()
95        {
96            // === GRAPH REPLAY ===
97            let exec = training_state.forward_graph_exec.as_ref().unwrap();
98            exec.launch(stream.raw())
99                .map_err(|e| eprintln!("[CUDA] Graph replay failed: {e}"))
100                .ok()?;
101            for _ in 0..cuda_blocks.len() {
102                input_is_a = !input_is_a;
103            }
104        } else {
105            // === Standard or first-capture forward ===
106            let capturing = use_graph && training_state.graph_cached_seq_len != seq_len;
107            if capturing {
108                // PMAT-063: Pre-allocate cuBLAS workspace before graph capture
109                if training_state.cublas_workspace.is_none() {
110                    training_state.cublas_workspace =
111                        super::super::gpu_backward_fallback::preallocate_cublas_workspace(trainer);
112                }
113                stream
114                    .begin_capture(CaptureMode::ThreadLocal)
115                    .map_err(|e| eprintln!("[CUDA] Graph capture begin failed: {e}"))
116                    .ok()?;
117            }
118
119            for (i, block) in cuda_blocks.iter_mut().enumerate() {
120                // SAFETY: ping-pong double-buffering. The two raw pointers reference distinct, non-overlapping device buffers (the `_a`/`_b` scratch pair); the boolean flag picks one as `&` input and the other as `&mut` output, so the resulting references never alias the same allocation.
121                let (gpu_input, gpu_output) = unsafe {
122                    if input_is_a {
123                        (&*scratch_a_ptr, &mut *scratch_b_ptr)
124                    } else {
125                        (&*scratch_b_ptr, &mut *scratch_a_ptr)
126                    }
127                };
128
129                // FALSIFY-CUDA-NF4-FORWARD-NAN-001: stream-ordered D2D copy.
130                // The sync `copy_from_buffer` (cuMemcpyDtoD) issues on the
131                // legacy default stream, which does NOT wait for this
132                // CU_STREAM_NON_BLOCKING stream — it can snapshot gpu_input
133                // before the previous layer finished writing it.
134                // SAFETY: both buffers outlive the stream-ordered copy
135                // (training_state and the ping-pong scratch live for the
136                // whole step); subsequent use is on the same stream.
137                unsafe {
138                    training_state.layer_inputs[i]
139                        .copy_from_buffer_async(gpu_input, stream)
140                        .map_err(|e| eprintln!("[CUDA] layer_input copy L{i}: {e}"))
141                        .ok()?;
142                }
143
144                // PMAT-483: Per-layer forward profiling
145                training_state.profiler_layer_start = Some(std::time::Instant::now());
146
147                if let Err(e) =
148                    block.forward(gpu_input, gpu_output, seq_len, stream, shared_scratch.as_mut())
149                {
150                    eprintln!(
151                        "[CUDA] Layer {i} forward failed: {e} (seq_len={seq_len} in={} out={} hidden={hidden_size})",
152                        gpu_input.len(), gpu_output.len(),
153                    );
154                    if capturing {
155                        let _ = stream.end_capture();
156                    }
157                    return None;
158                }
159
160                // PMAT-483: Record per-layer forward time
161                if let Some(start) = training_state.profiler_layer_start.take() {
162                    training_state.profiler_layer_fwd_us[i] = start.elapsed().as_micros() as u64;
163                }
164
165                input_is_a = !input_is_a;
166            }
167
168            if capturing {
169                match stream.end_capture() {
170                    Ok(graph) => match graph.instantiate() {
171                        Ok(exec) => {
172                            eprintln!(
173                                "[CUDA] Graph captured: {} layers, seq_len={seq_len}",
174                                cuda_blocks.len()
175                            );
176                            training_state.forward_graph_exec = Some(exec);
177                            training_state.graph_cached_seq_len = seq_len;
178                        }
179                        Err(e) => {
180                            eprintln!(
181                                "[CUDA] Graph instantiate failed: {e} — using non-graph path"
182                            );
183                        }
184                    },
185                    Err(e) => {
186                        eprintln!("[CUDA] Graph end_capture failed: {e} — using non-graph path");
187                    }
188                }
189            }
190        }
191
192        // SAFETY: ping-pong double-buffering. The two raw pointers reference distinct, non-overlapping device buffers (the `_a`/`_b` scratch pair); the boolean flag picks one as `&` input and the other as `&mut` output, so the resulting references never alias the same allocation.
193        let final_output = unsafe {
194            if input_is_a {
195                &*scratch_a_ptr
196            } else {
197                &*scratch_b_ptr
198            }
199        };
200
201        // Save blocks output for RMSNorm backward
202        if training_state.blocks_output.len() != final_output.len() {
203            training_state.blocks_output = trainer
204                .zeros(final_output.len())
205                .map_err(|e| eprintln!("[CUDA] blocks_output realloc failed: {e}"))
206                .ok()?;
207        }
208        // FALSIFY-CUDA-NF4-FORWARD-NAN-001: stream-ordered D2D copy (see above).
209        // SAFETY: both buffers outlive the stream-ordered copy; subsequent
210        // reads (RMSNorm backward) run on the same stream.
211        unsafe {
212            training_state
213                .blocks_output
214                .copy_from_buffer_async(final_output, stream)
215                .map_err(|e| eprintln!("[CUDA] blocks_output copy: {e}"))
216                .ok()?;
217        }
218
219        crate::autograd::cuda_backward::rms_norm_forward(
220            final_output,
221            &training_state.final_norm_weight,
222            &mut training_state.lm_head_hidden_buf,
223            seq_len as u32,
224            hidden_size as u32,
225            stream,
226        )
227        .map_err(|e| eprintln!("[CUDA] GPU RMSNorm forward failed: {e}"))
228        .ok()?;
229
230        crate::transformer::cuda_block::nan_scan_f32(
231            "final-norm",
232            &training_state.lm_head_hidden_buf,
233            seq_len * hidden_size,
234            stream,
235        );
236
237        Some(())
238    }
239    /// GPU-accelerated forward pass (inference-only, no layer input saving).
240    pub(super) fn forward_cuda_inference(
241        model: &Transformer,
242        token_ids: &[u32],
243        trainer: &CudaTrainer,
244        cuda_blocks: &mut [CudaBlock],
245        shared_scratch: &mut Option<CudaBlockScratch>,
246    ) -> Option<Vec<f32>> {
247        let seq_len = token_ids.len();
248        let hidden_size = model.config.hidden_size;
249
250        let hidden = model.embed_tokens.forward(token_ids);
251        let hidden_data = hidden.data();
252        let hidden_slice = hidden_data.as_slice().expect("contiguous hidden");
253
254        let mut gpu_input = trainer.upload(hidden_slice).ok()?;
255        let mut gpu_output = trainer.zeros(seq_len * hidden_size).ok()?;
256
257        let stream = trainer.stream();
258        for (i, block) in cuda_blocks.iter_mut().enumerate() {
259            if let Err(e) =
260                block.forward(&gpu_input, &mut gpu_output, seq_len, stream, shared_scratch.as_mut())
261            {
262                eprintln!("[CUDA] Layer {i} forward failed: {e}");
263                return None;
264            }
265            std::mem::swap(&mut gpu_input, &mut gpu_output);
266        }
267
268        if let Err(e) = stream.synchronize() {
269            eprintln!("[CUDA] Stream sync failed: {e}");
270            return None;
271        }
272
273        let result_data = trainer.download(&gpu_input).ok()?;
274        if result_data.iter().any(|v| !v.is_finite()) {
275            return None;
276        }
277
278        let result_tensor = crate::Tensor::from_vec(result_data, false);
279        let normed = model.norm.forward_batched(&result_tensor, seq_len, hidden_size);
280        let normed_data = normed.data();
281        let normed_slice = normed_data.as_slice().expect("contiguous normed");
282        Some(normed_slice.to_vec())
283    }
284    /// Forward pass dispatching to GPU. Returns logits as flat Vec<f32> [seq_len, vocab_size].
285    /// lm_head GEMM runs on GPU: hidden[seq, hidden] @ embed_T[hidden, vocab] -> logits[seq, vocab]
286    pub(super) fn forward_logits_gpu(&mut self, token_ids: &[u32]) -> Option<Vec<f32>> {
287        let seq_len = token_ids.len();
288        let vocab_size = self.model.config().vocab_size;
289        let hidden_size = self.model.config().hidden_size;
290
291        if self.gpu_training.is_some() {
292            let (trainer, blocks) = match (&self.cuda_trainer, &mut self.cuda_blocks) {
293                (Some(ref t), Some(ref mut b)) => (t, b),
294                _ => return None,
295            };
296            let mut training = self.gpu_training.take();
297            let result = Self::forward_cuda_training(
298                &self.model,
299                token_ids,
300                trainer,
301                blocks,
302                training.as_mut().expect("gpu_training was Some"),
303                &mut self.shared_scratch,
304            );
305            self.gpu_training = training;
306            result?;
307        } else {
308            let (trainer, blocks) = match (&self.cuda_trainer, &mut self.cuda_blocks) {
309                (Some(ref t), Some(ref mut b)) => (t, b),
310                _ => return None,
311            };
312            let normed_hidden = Self::forward_cuda_inference(
313                &self.model,
314                token_ids,
315                trainer,
316                blocks,
317                &mut self.shared_scratch,
318            )?;
319            let training = self.gpu_training.as_mut()?;
320            training
321                .lm_head_hidden_buf
322                .copy_from_host_at(&normed_hidden, 0)
323                .map_err(|e| eprintln!("[CUDA] lm_head forward: hidden upload failed: {e}"))
324                .ok()?;
325        }
326
327        let trainer = self.cuda_trainer.as_ref()?;
328        let training = self.gpu_training.as_mut()?;
329        let stream = trainer.stream();
330
331        eprintln!("[CUDA] lm_head BT: hidden_len={} embed_len={} logits_len={} seq={seq_len} h={hidden_size} v={vocab_size}",
332            training.lm_head_hidden_buf.len(), training.embed_original.len(), training.logits_buf.len());
333        if let Err(e) = crate::autograd::cuda_forward::gemm_forward_bt(
334            &training.lm_head_hidden_buf,
335            &training.embed_original,
336            &mut training.logits_buf,
337            seq_len as u32,
338            hidden_size as u32,
339            vocab_size as u32,
340            stream,
341        ) {
342            eprintln!("[CUDA] lm_head forward GEMM (BT) failed: {e}");
343            return None;
344        }
345
346        if let Err(e) = stream.synchronize() {
347            eprintln!("[CUDA] lm_head forward sync failed: {e}");
348            return None;
349        }
350
351        let full_logits = trainer
352            .download(&training.logits_buf)
353            .map_err(|e| eprintln!("[CUDA] lm_head forward: logits download failed: {e}"))
354            .ok()?;
355        Some(full_logits[..seq_len * vocab_size].to_vec())
356    }
357    /// PMAT-420: Inference forward + save layer inputs for backward.
358    /// Uses inference-style fresh buffers (no NaN) but saves layer inputs for GPU backward.
359    pub(super) fn forward_inference_saving_inputs(
360        &mut self,
361        token_ids: &[u32],
362    ) -> Option<Vec<f32>> {
363        let seq_len = token_ids.len();
364        let hidden_size = self.model.config().hidden_size;
365        let vocab_size = self.model.config().vocab_size;
366
367        let trainer = self.cuda_trainer.as_ref()?;
368        let blocks = self.cuda_blocks.as_mut()?;
369        let stream = trainer.stream();
370
371        let hidden = self.model.embed_tokens.forward(token_ids);
372        let hidden_data = hidden.data();
373        let hidden_slice = hidden_data.as_slice().expect("contiguous hidden");
374
375        let mut gpu_input = trainer.upload(hidden_slice).ok()?;
376        let mut gpu_output = trainer.zeros(seq_len * hidden_size).ok()?;
377
378        for (i, block) in blocks.iter_mut().enumerate() {
379            if let Some(ref mut training) = self.gpu_training {
380                if i < training.layer_inputs.len() {
381                    if training.layer_inputs[i].len() != gpu_input.len() {
382                        if let Ok(buf) = trainer.zeros(gpu_input.len()) {
383                            training.layer_inputs[i] = buf;
384                        }
385                    }
386                    training.layer_inputs[i]
387                        .copy_from_buffer(&gpu_input)
388                        .map_err(|e| eprintln!("[CUDA] layer_input copy L{i}: {e}"))
389                        .ok();
390                }
391            }
392
393            if let Err(e) = block.forward(
394                &gpu_input,
395                &mut gpu_output,
396                seq_len,
397                stream,
398                self.shared_scratch.as_mut(),
399            ) {
400                eprintln!("[CUDA] Layer {i} forward failed: {e}");
401                return None;
402            }
403            std::mem::swap(&mut gpu_input, &mut gpu_output);
404        }
405
406        stream.synchronize().ok()?;
407
408        // Save blocks_output for RMSNorm backward
409        if let Some(ref mut training) = self.gpu_training {
410            if training.blocks_output.len() != gpu_input.len() {
411                if let Ok(buf) = trainer.zeros(gpu_input.len()) {
412                    training.blocks_output = buf;
413                }
414            }
415            training
416                .blocks_output
417                .copy_from_buffer(&gpu_input)
418                .map_err(|e| eprintln!("[CUDA] blocks_output copy: {e}"))
419                .ok();
420        }
421
422        let result = trainer.download(&gpu_input).ok()?;
423        if result.iter().any(|v| !v.is_finite()) {
424            eprintln!("[CUDA] NaN in forward output — inference-style forward failed");
425            return None;
426        }
427
428        // CPU RMSNorm
429        let result_tensor = crate::autograd::Tensor::from_vec(result, false);
430        let normed = self.model.norm.forward_batched(&result_tensor, seq_len, hidden_size);
431        let normed_data = normed.data();
432        let normed_slice = normed_data.as_slice().expect("contiguous normed");
433
434        // Save normed hidden for lm_head backward
435        if let Some(ref mut training) = self.gpu_training {
436            if let Ok(buf) = trainer.upload(normed_slice) {
437                training.lm_head_hidden_buf = buf;
438            }
439        }
440
441        // CPU lm_head
442        let lm_weight = self.model.lm_head.as_ref().unwrap_or(&self.model.embed_tokens.weight);
443        let lm_data = lm_weight.data();
444        let lm_slice = lm_data.as_slice().expect("contiguous lm_head");
445        let logits = crate::autograd::ops::matmul::matmul_nt_compute(
446            normed_slice,
447            lm_slice,
448            seq_len,
449            hidden_size,
450            vocab_size,
451        );
452        Some(logits)
453    }
454    /// GPU forward with logits staying GPU-resident (KAIZEN-064).
455    /// After this call, `training.logits_buf` contains logits on GPU. Returns true on success.
456    pub(super) fn forward_logits_gpu_resident(&mut self, token_ids: &[u32]) -> bool {
457        let seq_len = token_ids.len();
458        let vocab_size = self.model.config().vocab_size;
459        let hidden_size = self.model.config().hidden_size;
460
461        if self.gpu_training.is_some() {
462            let (trainer, blocks) = match (&self.cuda_trainer, &mut self.cuda_blocks) {
463                (Some(ref t), Some(ref mut b)) => (t, b),
464                _ => {
465                    eprintln!("[RES-FALSE] no trainer/blocks");
466                    return false;
467                }
468            };
469            let mut training = self.gpu_training.take();
470            let result = Self::forward_cuda_training(
471                &self.model,
472                token_ids,
473                trainer,
474                blocks,
475                training.as_mut().expect("gpu_training was Some"),
476                &mut self.shared_scratch,
477            );
478            self.gpu_training = training;
479            if result.is_none() {
480                eprintln!("[RES-FALSE] forward_cuda_training returned None");
481                return false;
482            }
483        } else {
484            let (trainer, blocks) = match (&self.cuda_trainer, &mut self.cuda_blocks) {
485                (Some(ref t), Some(ref mut b)) => (t, b),
486                _ => return false,
487            };
488            let normed_hidden = match Self::forward_cuda_inference(
489                &self.model,
490                token_ids,
491                trainer,
492                blocks,
493                &mut self.shared_scratch,
494            ) {
495                Some(h) => h,
496                None => return false,
497            };
498            let training = match self.gpu_training.as_mut() {
499                Some(t) => t,
500                None => return false,
501            };
502            if training.lm_head_hidden_buf.copy_from_host_at(&normed_hidden, 0).is_err() {
503                eprintln!("[CUDA] lm_head forward: hidden upload failed");
504                return false;
505            }
506        }
507
508        let (trainer, training) = match (&self.cuda_trainer, &mut self.gpu_training) {
509            (Some(ref t), Some(ref mut tr)) => (t, tr),
510            _ => {
511                eprintln!("[RES-FALSE] no trainer/training");
512                return false;
513            }
514        };
515
516        let stream = trainer.stream();
517
518        if crate::autograd::cuda_forward::gemm_forward_bt(
519            &training.lm_head_hidden_buf,
520            &training.embed_original,
521            &mut training.logits_buf,
522            seq_len as u32,
523            hidden_size as u32,
524            vocab_size as u32,
525            stream,
526        )
527        .is_err()
528        {
529            eprintln!("[CUDA] lm_head forward GEMM (BT) failed");
530            eprintln!("[RES-FALSE] BT GEMM failed");
531            return false;
532        }
533
534        crate::transformer::cuda_block::nan_scan_f32(
535            "lm-head-logits",
536            &training.logits_buf,
537            seq_len * vocab_size,
538            stream,
539        );
540
541        true
542    }
543}
entrenar/finetune/instruct_pipeline/cuda_forward.rs

entrenar/finetune/instruct_pipeline/
cuda_forward.rs
