realizar 0.8.5

Pure Rust ML inference engine built from scratch - model serving for GGUF and safetensors
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145

impl CudaExecutor {
    /// Validate workspace is initialized and large enough for batch size M.
    pub(crate) fn validate_batched_workspace(&self, m: u32, positions_len: usize) -> Result<(), GpuError> {
        if !self.workspace.initialized {
            return Err(GpuError::InvalidLaunchConfig(
                "PAR-111: Workspace not initialized".to_string(),
            ));
        }
        if self.workspace.buffer_capacity < m as usize {
            return Err(GpuError::InvalidLaunchConfig(format!(
                "PAR-111: Workspace buffer_capacity {} < m {}",
                self.workspace.buffer_capacity, m
            )));
        }
        if positions_len != m as usize {
            return Err(GpuError::InvalidLaunchConfig(format!(
                "PAR-111: positions.len() {} != m {}",
                positions_len, m
            )));
        }
        Ok(())
    }

    /// PAR-111: Batched forward pass for M sequences through a single layer.
    ///
    /// Processes M sequences in parallel using batched GEMV kernels.
    /// Each sequence has independent KV cache state.
    ///
    /// # Performance Benefit
    ///
    /// Batched GEMV reads/dequantizes weights ONCE for all M inputs:
    /// - M=1: Baseline throughput (~360 tok/s)
    /// - M=4: 16x GEMV speedup → 857+ tok/s aggregate
    #[allow(clippy::too_many_arguments, clippy::too_many_lines)]
    pub fn transformer_layer_batched(
        &mut self,
        input: &GpuBuffer<f32>,
        layer_idx: usize,
        layer_weights: &ValidatedLayerWeights,
        m: u32,
        positions: &[u32],
        hidden_dim: u32,
        intermediate_dim: u32,
        epsilon: f32,
    ) -> Result<(), GpuError> {
        self.validate_batched_workspace(m, positions.len())?;

        let q_dim = (self.kv_num_heads * self.kv_head_dim) as u32;
        let kv_dim = (self.kv_num_kv_heads * self.kv_head_dim) as u32;

        // Get batched buffer pointers (M× larger buffers allocated by init_batched_workspace)
        let hidden_buf1_ptr = self.workspace.hidden_buf1.as_ref()
            .ok_or_else(|| GpuError::InvalidLaunchConfig("PAR-111: hidden_buf1 not initialized".to_string()))?.as_ptr();
        let hidden_buf1_len = self.workspace.hidden_buf1.as_ref()
            .ok_or_else(|| GpuError::InvalidLaunchConfig("PAR-111: hidden_buf1 not initialized".to_string()))?.len();
        let hidden_buf2_ptr = self.workspace.hidden_buf2.as_ref()
            .ok_or_else(|| GpuError::InvalidLaunchConfig("PAR-111: hidden_buf2 not initialized".to_string()))?.as_ptr();
        let hidden_buf2_len = self.workspace.hidden_buf2.as_ref()
            .ok_or_else(|| GpuError::InvalidLaunchConfig("PAR-111: hidden_buf2 not initialized".to_string()))?.len();
        let input_staging_ptr = self.workspace.input_staging.as_ref()
            .ok_or_else(|| GpuError::InvalidLaunchConfig("PAR-111: input_staging not initialized".to_string()))?.as_ptr();
        let input_staging_len = self.workspace.input_staging.as_ref()
            .ok_or_else(|| GpuError::InvalidLaunchConfig("PAR-111: input_staging not initialized".to_string()))?.len();
        let q_buf_ptr = self.workspace.q_buf.as_ref()
            .ok_or_else(|| GpuError::InvalidLaunchConfig("PAR-111: q_buf not initialized".to_string()))?.as_ptr();
        let q_buf_len = self.workspace.q_buf.as_ref()
            .ok_or_else(|| GpuError::InvalidLaunchConfig("PAR-111: q_buf not initialized".to_string()))?.len();
        let k_buf_ptr = self.workspace.k_buf.as_ref()
            .ok_or_else(|| GpuError::InvalidLaunchConfig("PAR-111: k_buf not initialized".to_string()))?.as_ptr();
        let k_buf_len = self.workspace.k_buf.as_ref()
            .ok_or_else(|| GpuError::InvalidLaunchConfig("PAR-111: k_buf not initialized".to_string()))?.len();
        let v_buf_ptr = self.workspace.v_buf.as_ref()
            .ok_or_else(|| GpuError::InvalidLaunchConfig("PAR-111: v_buf not initialized".to_string()))?.as_ptr();
        let v_buf_len = self.workspace.v_buf.as_ref()
            .ok_or_else(|| GpuError::InvalidLaunchConfig("PAR-111: v_buf not initialized".to_string()))?.len();
        let ffn_gate_ptr = self.workspace.ffn_gate_buf.as_ref()
            .ok_or_else(|| GpuError::InvalidLaunchConfig("PAR-111: ffn_gate_buf not initialized".to_string()))?.as_ptr();
        let ffn_gate_len = self.workspace.ffn_gate_buf.as_ref()
            .ok_or_else(|| GpuError::InvalidLaunchConfig("PAR-111: ffn_gate_buf not initialized".to_string()))?.len();
        let ffn_up_ptr = self.workspace.ffn_up_buf.as_ref()
            .ok_or_else(|| GpuError::InvalidLaunchConfig("PAR-111: ffn_up_buf not initialized".to_string()))?.as_ptr();
        let ffn_up_len = self.workspace.ffn_up_buf.as_ref()
            .ok_or_else(|| GpuError::InvalidLaunchConfig("PAR-111: ffn_up_buf not initialized".to_string()))?.len();
        let ffn_act_ptr = self.workspace.ffn_act_buf.as_ref()
            .ok_or_else(|| GpuError::InvalidLaunchConfig("PAR-111: ffn_act_buf not initialized".to_string()))?.as_ptr();
        let ffn_act_len = self.workspace.ffn_act_buf.as_ref()
            .ok_or_else(|| GpuError::InvalidLaunchConfig("PAR-111: ffn_act_buf not initialized".to_string()))?.len();
        let attn_out_ptr = self.workspace.attn_out_buf.as_ref()
            .ok_or_else(|| GpuError::InvalidLaunchConfig("PAR-111: attn_out_buf not initialized".to_string()))?.as_ptr();
        let attn_out_len = self.workspace.attn_out_buf.as_ref()
            .ok_or_else(|| GpuError::InvalidLaunchConfig("PAR-111: attn_out_buf not initialized".to_string()))?.len();

        // Create temporary buffer wrappers (M× sized)
        // SAFETY: Pointers valid from workspace allocation, lengths verified
        let hidden_buf1 = unsafe { GpuBuffer::<f32>::from_raw_parts(hidden_buf1_ptr, hidden_buf1_len) };
        let hidden_buf2 = unsafe { GpuBuffer::<f32>::from_raw_parts(hidden_buf2_ptr, hidden_buf2_len) };
        let input_staging = unsafe { GpuBuffer::<f32>::from_raw_parts(input_staging_ptr, input_staging_len) };
        let q_buf = unsafe { GpuBuffer::<f32>::from_raw_parts(q_buf_ptr, q_buf_len) };
        let k_buf = unsafe { GpuBuffer::<f32>::from_raw_parts(k_buf_ptr, k_buf_len) };
        // SAFETY: v_buf_ptr valid from workspace allocation, v_buf_len verified
        let v_buf = unsafe { GpuBuffer::<f32>::from_raw_parts(v_buf_ptr, v_buf_len) };
        // SAFETY: pointers valid from workspace allocation above, lengths verified at allocation
        let ffn_gate_buf = unsafe { GpuBuffer::<f32>::from_raw_parts(ffn_gate_ptr, ffn_gate_len) };
        let ffn_up_buf = unsafe { GpuBuffer::<f32>::from_raw_parts(ffn_up_ptr, ffn_up_len) };
        let ffn_act_buf = unsafe { GpuBuffer::<f32>::from_raw_parts(ffn_act_ptr, ffn_act_len) };
        let attn_out_buf = unsafe { GpuBuffer::<f32>::from_raw_parts(attn_out_ptr, attn_out_len) };

        // Phase 1: RMSNorm + QKV projections + bias + RoPE
        self.batched_qkv_rope_phase(
            input, &hidden_buf1, &q_buf, &k_buf, &v_buf,
            q_buf_ptr, k_buf_ptr, v_buf_ptr, hidden_buf1_ptr,
            layer_idx, layer_weights, m, positions,
            hidden_dim, q_dim, kv_dim, epsilon,
        )?;

        // Phase 2: Attention + output projection + residuals + FFN
        self.batched_attn_ffn_phase(
            input, &hidden_buf1, &hidden_buf2, &input_staging,
            &q_buf, &k_buf, &v_buf, &attn_out_buf,
            &ffn_gate_buf, &ffn_up_buf, &ffn_act_buf,
            q_buf_ptr, k_buf_ptr, v_buf_ptr,
            attn_out_ptr, hidden_buf1_ptr, ffn_gate_ptr, ffn_up_ptr, ffn_act_ptr,
            layer_idx, layer_weights, m, positions,
            hidden_dim, intermediate_dim, q_dim, kv_dim, epsilon,
        )?;

        // Prevent Drop from freeing borrowed memory
        std::mem::forget(hidden_buf1);
        std::mem::forget(hidden_buf2);
        std::mem::forget(input_staging);
        std::mem::forget(q_buf);
        std::mem::forget(k_buf);
        std::mem::forget(v_buf);
        std::mem::forget(attn_out_buf);
        std::mem::forget(ffn_gate_buf);
        std::mem::forget(ffn_up_buf);
        std::mem::forget(ffn_act_buf);

        // Output is in hidden_buf2 (M × hidden_dim packed)
        Ok(())
    }
}

include!("batched_qkv.rs");
include!("batched_ffn.rs");