/home/noah/src/realizar/src/brick/mod.rs
Line | Count | Source |
1 | | //! ComputeBrick: Token-centric, self-verifying compute units. |
2 | | //! |
3 | | //! Per CBTOP-SPEC-001 and SHOWCASE-BRICK-001, every compute operation is a brick with: |
4 | | //! - Token budget (tok/sec performance target) |
5 | | //! - Assertions (falsifiable correctness claims) |
6 | | //! - Verification (self-checking via baseline comparison) |
7 | | //! |
8 | | //! # Toyota Way Integration |
9 | | //! - **Jidoka**: Stop-the-line on budget violation |
10 | | //! - **Poka-Yoke**: Type-safe brick composition |
11 | | //! - **Genchi Genbutsu**: Real metrics from hardware |
12 | | //! - **Mieruka**: Visual control via cbtop TUI |
13 | | //! |
14 | | //! # References |
15 | | //! - Popper, K. (1959). "The Logic of Scientific Discovery." |
16 | | //! - Ohno, T. (1988). "Toyota Production System." |
17 | | //! - Little, J.D.C. (1961). "A Proof for the Queuing Formula: L = λW." |
18 | | |
19 | | use std::fmt; |
20 | | use std::time::Instant; |
21 | | |
22 | | use crate::quantize::Q8_0Block; |
23 | | |
24 | | // PMAT-802: Extracted modules |
25 | | #[cfg(feature = "cuda")] |
26 | | mod fused; |
27 | | #[cfg(feature = "cuda")] |
28 | | pub use fused::{CoalescedDp4aBrick, FusedFfnBrick}; |
29 | | |
30 | | // Phase 14: BrickTracer for GPU/CPU parity debugging |
31 | | pub mod tracer; |
32 | | pub use tracer::{BrickTracer, TraceEvent, TraceComparison, TraceDiff}; |
33 | | |
34 | | // ============================================================================ |
35 | | // Core Types |
36 | | // ============================================================================ |
37 | | |
38 | | /// Performance budget expressed in token terms. |
39 | | /// Aligns compute costs with LLM inference metrics. |
40 | | /// |
41 | | /// Per Little's Law (1961): throughput = 1 / latency |
42 | | #[derive(Debug, Clone, Copy, PartialEq)] |
43 | | pub struct TokenBudget { |
44 | | /// Latency budget per token (microseconds) |
45 | | pub us_per_token: f64, |
46 | | /// Throughput target (tokens/second) |
47 | | pub tokens_per_sec: f64, |
48 | | /// Batch size for amortization |
49 | | pub batch_size: usize, |
50 | | } |
51 | | |
52 | | impl TokenBudget { |
53 | | /// Create budget from latency target. |
54 | | /// 50µs/token = 20,000 tokens/sec |
55 | | #[must_use] |
56 | 385 | pub fn from_latency(us_per_token: f64) -> Self { |
57 | 385 | Self { |
58 | 385 | us_per_token, |
59 | 385 | tokens_per_sec: 1_000_000.0 / us_per_token, |
60 | 385 | batch_size: 1, |
61 | 385 | } |
62 | 385 | } |
63 | | |
64 | | /// Create budget from throughput target. |
65 | | /// 20,000 tokens/sec = 50µs/token |
66 | | #[must_use] |
67 | 4 | pub fn from_throughput(tokens_per_sec: f64) -> Self { |
68 | 4 | Self { |
69 | 4 | us_per_token: 1_000_000.0 / tokens_per_sec, |
70 | 4 | tokens_per_sec, |
71 | 4 | batch_size: 1, |
72 | 4 | } |
73 | 4 | } |
74 | | |
75 | | /// Create budget with batch size. |
76 | | #[must_use] |
77 | 1 | pub fn with_batch_size(mut self, batch_size: usize) -> Self { |
78 | 1 | self.batch_size = batch_size; |
79 | 1 | self |
80 | 1 | } |
81 | | |
82 | | /// Check if actual performance meets budget. |
83 | | #[must_use] |
84 | 65 | pub fn is_met(&self, actual_us_per_token: f64) -> bool { |
85 | 65 | actual_us_per_token <= self.us_per_token |
86 | 65 | } |
87 | | |
88 | | /// Calculate gap factor (actual / budget). |
89 | | /// < 1.0 means under budget (good), > 1.0 means over budget (bad). |
90 | | #[must_use] |
91 | 9 | pub fn gap_factor(&self, actual_us_per_token: f64) -> f64 { |
92 | 9 | actual_us_per_token / self.us_per_token |
93 | 9 | } |
94 | | } |
95 | | |
96 | | impl Default for TokenBudget { |
97 | 1 | fn default() -> Self { |
98 | 1 | Self::from_latency(100.0) // 100µs = 10k tok/s default |
99 | 1 | } |
100 | | } |
101 | | |
102 | | /// Result of ComputeBrick execution with token metrics. |
103 | | #[derive(Debug, Clone)] |
104 | | pub struct TokenResult<T> { |
105 | | /// Computed output |
106 | | pub output: T, |
107 | | /// Number of tokens processed |
108 | | pub tokens_processed: usize, |
109 | | /// Actual latency (microseconds/token) |
110 | | pub us_per_token: f64, |
111 | | /// Actual throughput (tokens/second) |
112 | | pub tokens_per_sec: f64, |
113 | | /// Did we meet the budget? |
114 | | pub budget_met: bool, |
115 | | } |
116 | | |
117 | | impl<T: Default> Default for TokenResult<T> { |
118 | 0 | fn default() -> Self { |
119 | 0 | Self { |
120 | 0 | output: T::default(), |
121 | 0 | tokens_processed: 0, |
122 | 0 | us_per_token: 0.0, |
123 | 0 | tokens_per_sec: 0.0, |
124 | 0 | budget_met: true, |
125 | 0 | } |
126 | 0 | } |
127 | | } |
128 | | |
129 | | impl<T> TokenResult<T> { |
130 | | /// Create a new token result. |
131 | 62 | pub fn new(output: T, tokens: usize, elapsed_us: f64, budget: &TokenBudget) -> Self { |
132 | 62 | let us_per_token = elapsed_us / tokens.max(1) as f64; |
133 | 62 | let tokens_per_sec = if us_per_token > 0.0 { |
134 | 60 | 1_000_000.0 / us_per_token |
135 | | } else { |
136 | 2 | 0.0 |
137 | | }; |
138 | | |
139 | 62 | Self { |
140 | 62 | output, |
141 | 62 | tokens_processed: tokens, |
142 | 62 | us_per_token, |
143 | 62 | tokens_per_sec, |
144 | 62 | budget_met: budget.is_met(us_per_token), |
145 | 62 | } |
146 | 62 | } |
147 | | } |
148 | | |
149 | | /// Errors from ComputeBrick execution. |
150 | | /// Tells you exactly what failed (Jidoka: stop and signal). |
151 | | #[derive(Debug)] |
152 | | pub enum BrickError { |
153 | | /// Assertion failed during verification. |
154 | | AssertionFailed { |
155 | | /// Assertion name |
156 | | name: String, |
157 | | /// Expected value |
158 | | expected: String, |
159 | | /// Actual value |
160 | | actual: String, |
161 | | }, |
162 | | /// Performance budget exceeded. |
163 | | BudgetExceeded { |
164 | | /// Budget limit in µs |
165 | | limit_us: f64, |
166 | | /// Actual time in µs |
167 | | actual_us: f64, |
168 | | }, |
169 | | /// Compute operation failed. |
170 | | ComputeError(String), |
171 | | /// Invalid input. |
172 | | InvalidInput(String), |
173 | | } |
174 | | |
175 | | impl fmt::Display for BrickError { |
176 | 6 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
177 | 6 | match self { |
178 | | Self::AssertionFailed { |
179 | 1 | name, |
180 | 1 | expected, |
181 | 1 | actual, |
182 | | } => { |
183 | 1 | write!( |
184 | 1 | f, |
185 | 1 | "Assertion failed: {name} - expected {expected}, got {actual}" |
186 | | ) |
187 | | }, |
188 | | Self::BudgetExceeded { |
189 | 2 | limit_us, |
190 | 2 | actual_us, |
191 | | } => { |
192 | 2 | write!( |
193 | 2 | f, |
194 | 2 | "Budget exceeded: {limit_us:.1}µs/tok limit, {actual_us:.1}µs/tok actual" |
195 | | ) |
196 | | }, |
197 | 1 | Self::ComputeError(msg) => write!(f, "Compute error: {msg}"), |
198 | 2 | Self::InvalidInput(msg) => write!(f, "Invalid input: {msg}"), |
199 | | } |
200 | 6 | } |
201 | | } |
202 | | |
203 | | impl std::error::Error for BrickError {} |
204 | | |
205 | | // ============================================================================ |
206 | | // Brick Assertions |
207 | | // ============================================================================ |
208 | | |
209 | | /// Falsifiable assertion for a brick (Popper criterion). |
210 | | #[derive(Debug, Clone)] |
211 | | pub struct BrickAssertion { |
212 | | /// Assertion name (for error messages) |
213 | | pub name: String, |
214 | | /// Description of what is being asserted |
215 | | pub description: String, |
216 | | /// Assertion type |
217 | | pub kind: AssertionKind, |
218 | | } |
219 | | |
220 | | /// Types of assertions a brick can make. |
221 | | #[derive(Debug, Clone)] |
222 | | pub enum AssertionKind { |
223 | | /// Output matches scalar baseline within tolerance. |
224 | | EquivScalar { |
225 | | /// Tolerance for comparison |
226 | | tolerance: f64, |
227 | | }, |
228 | | /// Output contains no NaN values. |
229 | | NoNaN, |
230 | | /// Output contains no Inf values. |
231 | | NoInf, |
232 | | /// Output values within bounds. |
233 | | Bounds { |
234 | | /// Minimum allowed value |
235 | | min: f64, |
236 | | /// Maximum allowed value |
237 | | max: f64, |
238 | | }, |
239 | | /// Budget must be met. |
240 | | BudgetMet, |
241 | | /// Custom assertion (returns true if passed). |
242 | | Custom { |
243 | | /// Name of the custom check |
244 | | check_name: String, |
245 | | }, |
246 | | } |
247 | | |
248 | | impl BrickAssertion { |
249 | | /// Create assertion that output matches scalar baseline. |
250 | 2 | pub fn equiv_scalar(tolerance: f64) -> Self { |
251 | 2 | Self { |
252 | 2 | name: "equiv_scalar".to_string(), |
253 | 2 | description: format!("Output matches scalar baseline within {tolerance}"), |
254 | 2 | kind: AssertionKind::EquivScalar { tolerance }, |
255 | 2 | } |
256 | 2 | } |
257 | | |
258 | | /// Create assertion that output contains no NaN. |
259 | 73 | pub fn no_nan() -> Self { |
260 | 73 | Self { |
261 | 73 | name: "no_nan".to_string(), |
262 | 73 | description: "Output contains no NaN values".to_string(), |
263 | 73 | kind: AssertionKind::NoNaN, |
264 | 73 | } |
265 | 73 | } |
266 | | |
267 | | /// Create assertion that output contains no Inf. |
268 | 72 | pub fn no_inf() -> Self { |
269 | 72 | Self { |
270 | 72 | name: "no_inf".to_string(), |
271 | 72 | description: "Output contains no Inf values".to_string(), |
272 | 72 | kind: AssertionKind::NoInf, |
273 | 72 | } |
274 | 72 | } |
275 | | |
276 | | /// Create assertion that values are within bounds. |
277 | 5 | pub fn bounds(min: f64, max: f64) -> Self { |
278 | 5 | Self { |
279 | 5 | name: "bounds".to_string(), |
280 | 5 | description: format!("Output values in [{min}, {max}]"), |
281 | 5 | kind: AssertionKind::Bounds { min, max }, |
282 | 5 | } |
283 | 5 | } |
284 | | |
285 | | /// Create assertion that budget is met. |
286 | 73 | pub fn budget_met() -> Self { |
287 | 73 | Self { |
288 | 73 | name: "budget_met".to_string(), |
289 | 73 | description: "Performance budget is met".to_string(), |
290 | 73 | kind: AssertionKind::BudgetMet, |
291 | 73 | } |
292 | 73 | } |
293 | | |
294 | | /// Check assertion against f32 slice output. |
295 | 183 | pub fn check_f32(&self, output: &[f32], budget_met: bool) -> Result<(), BrickError> { |
296 | 183 | match &self.kind { |
297 | | AssertionKind::NoNaN => { |
298 | 72.7k | if let Some(idx1 ) = output.iter()61 .position61 (|x| x.is_nan()) { |
299 | 1 | return Err(BrickError::AssertionFailed { |
300 | 1 | name: self.name.clone(), |
301 | 1 | expected: "no NaN".to_string(), |
302 | 1 | actual: format!("NaN at index {idx}"), |
303 | 1 | }); |
304 | 60 | } |
305 | | }, |
306 | | AssertionKind::NoInf => { |
307 | 72.7k | if let Some(idx0 ) = output.iter()59 .position59 (|x| x.is_infinite()) { |
308 | 0 | return Err(BrickError::AssertionFailed { |
309 | 0 | name: self.name.clone(), |
310 | 0 | expected: "no Inf".to_string(), |
311 | 0 | actual: format!("Inf at index {idx}"), |
312 | 0 | }); |
313 | 59 | } |
314 | | }, |
315 | 3 | AssertionKind::Bounds { min, max } => { |
316 | 9 | for (idx, &val) in output3 .iter3 ().enumerate3 () { |
317 | 9 | if (val as f64) < *min || (val as f64) > *max { |
318 | 1 | return Err(BrickError::AssertionFailed { |
319 | 1 | name: self.name.clone(), |
320 | 1 | expected: format!("value in [{min}, {max}]"), |
321 | 1 | actual: format!("value {val} at index {idx}"), |
322 | 1 | }); |
323 | 8 | } |
324 | | } |
325 | | }, |
326 | | AssertionKind::BudgetMet => { |
327 | 58 | if !budget_met { |
328 | 0 | return Err(BrickError::AssertionFailed { |
329 | 0 | name: self.name.clone(), |
330 | 0 | expected: "budget met".to_string(), |
331 | 0 | actual: "budget exceeded".to_string(), |
332 | 0 | }); |
333 | 58 | } |
334 | | }, |
335 | 2 | AssertionKind::EquivScalar { .. } | AssertionKind::Custom { .. } => { |
336 | 2 | // These require external comparison, skip for basic check |
337 | 2 | }, |
338 | | } |
339 | 181 | Ok(()) |
340 | 183 | } |
341 | | } |
342 | | |
343 | | // ============================================================================ |
344 | | // Brick Verification |
345 | | // ============================================================================ |
346 | | |
347 | | /// Result of brick verification. |
348 | | #[derive(Debug, Clone)] |
349 | | pub struct BrickVerification { |
350 | | /// All assertions passed? |
351 | | pub is_valid: bool, |
352 | | /// Individual assertion results |
353 | | pub results: Vec<(String, bool, String)>, |
354 | | } |
355 | | |
356 | | impl BrickVerification { |
357 | | /// Create a passing verification. |
358 | 10 | pub fn pass() -> Self { |
359 | 10 | Self { |
360 | 10 | is_valid: true, |
361 | 10 | results: vec![], |
362 | 10 | } |
363 | 10 | } |
364 | | |
365 | | /// Create a failing verification. |
366 | 1 | pub fn fail(name: &str, reason: &str) -> Self { |
367 | 1 | Self { |
368 | 1 | is_valid: false, |
369 | 1 | results: vec![(name.to_string(), false, reason.to_string())], |
370 | 1 | } |
371 | 1 | } |
372 | | |
373 | | /// Add an assertion result. |
374 | 2 | pub fn add(&mut self, name: &str, passed: bool, message: &str) { |
375 | 2 | self.results |
376 | 2 | .push((name.to_string(), passed, message.to_string())); |
377 | 2 | if !passed { |
378 | 1 | self.is_valid = false; |
379 | 1 | } |
380 | 2 | } |
381 | | } |
382 | | |
383 | | // ============================================================================ |
384 | | // ComputeBrick Trait |
385 | | // ============================================================================ |
386 | | |
387 | | /// Core trait for self-verifying, token-centric compute units. |
388 | | /// |
389 | | /// Every brick must: |
390 | | /// 1. Have at least one assertion (Popper criterion) |
391 | | /// 2. Have a non-zero budget (accountability) |
392 | | /// 3. Be verifiable against baseline |
393 | | /// |
394 | | /// # Invariants (PROBAR-SPEC-009 §3) |
395 | | /// - `assertions().len() > 0` (at least one falsifiable claim) |
396 | | /// - `budget().us_per_token > 0` (performance accountability) |
397 | | /// - `verify()` checks ALL assertions |
398 | | pub trait ComputeBrick: Send + Sync { |
399 | | /// Output type of this brick. |
400 | | type Output; |
401 | | |
402 | | /// Brick name for identification. |
403 | | fn name(&self) -> &'static str; |
404 | | |
405 | | /// Token throughput budget. |
406 | | fn budget(&self) -> TokenBudget; |
407 | | |
408 | | /// Falsifiable assertions for this brick. |
409 | | fn assertions(&self) -> Vec<BrickAssertion>; |
410 | | |
411 | | /// Verify brick state without running. |
412 | 7 | fn verify(&self) -> BrickVerification { |
413 | 7 | let assertions = self.assertions(); |
414 | 7 | if assertions.is_empty() { |
415 | 0 | return BrickVerification::fail( |
416 | 0 | self.name(), |
417 | 0 | "No assertions defined (Popper violation)", |
418 | | ); |
419 | 7 | } |
420 | | |
421 | 7 | let budget = self.budget(); |
422 | 7 | if budget.us_per_token <= 0.0 { |
423 | 0 | return BrickVerification::fail(self.name(), "Zero or negative budget"); |
424 | 7 | } |
425 | | |
426 | 7 | BrickVerification::pass() |
427 | 7 | } |
428 | | |
429 | | /// Can this brick run? (Jidoka gate) |
430 | 5 | fn can_run(&self) -> bool { |
431 | 5 | self.verify().is_valid |
432 | 5 | } |
433 | | } |
434 | | |
435 | | // ============================================================================ |
436 | | // Transformer Brick Implementations |
437 | | // ============================================================================ |
438 | | |
439 | | /// RMSNorm brick - layer normalization. |
440 | | #[derive(Debug)] |
441 | | pub struct RmsNormBrick { |
442 | | /// Weight vector |
443 | | pub weight: Vec<f32>, |
444 | | /// Epsilon for numerical stability |
445 | | pub eps: f32, |
446 | | /// Budget |
447 | | budget: TokenBudget, |
448 | | } |
449 | | |
450 | | impl RmsNormBrick { |
451 | | /// Create a new RMSNorm brick. |
452 | 113 | pub fn new(weight: Vec<f32>, eps: f32) -> Self { |
453 | 113 | Self { |
454 | 113 | weight, |
455 | 113 | eps, |
456 | 113 | budget: TokenBudget::from_latency(1.5), // 1.5µs target |
457 | 113 | } |
458 | 113 | } |
459 | | |
460 | | /// Set custom budget. |
461 | | #[must_use] |
462 | 4 | pub fn with_budget(mut self, budget: TokenBudget) -> Self { |
463 | 4 | self.budget = budget; |
464 | 4 | self |
465 | 4 | } |
466 | | |
467 | | /// Run RMSNorm on input. |
468 | 58 | pub fn run(&self, input: &[f32]) -> Result<TokenResult<Vec<f32>>, BrickError> { |
469 | 58 | if input.len() != self.weight.len() { |
470 | 0 | return Err(BrickError::InvalidInput(format!( |
471 | 0 | "Input len {} != weight len {}", |
472 | 0 | input.len(), |
473 | 0 | self.weight.len() |
474 | 0 | ))); |
475 | 58 | } |
476 | | |
477 | 58 | let start = Instant::now(); |
478 | | |
479 | | // Compute RMS |
480 | 72.7k | let rms58 = (input58 .iter58 ().map58 (|x| x * x).sum58 ::<f32>() / input.len() as f3258 + self.eps58 ).sqrt58 (); |
481 | | |
482 | | // Normalize and scale |
483 | 58 | let output: Vec<f32> = input |
484 | 58 | .iter() |
485 | 58 | .zip(self.weight.iter()) |
486 | 72.7k | .map58 (|(x, w)| (x / rms) * w) |
487 | 58 | .collect(); |
488 | | |
489 | 58 | let elapsed_us = start.elapsed().as_micros() as f64; |
490 | 58 | let result = TokenResult::new(output, 1, elapsed_us, &self.budget); |
491 | | |
492 | | // Check assertions |
493 | 174 | for assertion in self58 .assertions58 () { |
494 | 174 | assertion.check_f32(&result.output, result.budget_met)?0 ; |
495 | | } |
496 | | |
497 | 58 | Ok(result) |
498 | 58 | } |
499 | | } |
500 | | |
501 | | impl ComputeBrick for RmsNormBrick { |
502 | | type Output = Vec<f32>; |
503 | | |
504 | 2 | fn name(&self) -> &'static str { |
505 | 2 | "rms_norm" |
506 | 2 | } |
507 | | |
508 | 16 | fn budget(&self) -> TokenBudget { |
509 | 16 | self.budget |
510 | 16 | } |
511 | | |
512 | 65 | fn assertions(&self) -> Vec<BrickAssertion> { |
513 | 65 | vec![ |
514 | 65 | BrickAssertion::no_nan(), |
515 | 65 | BrickAssertion::no_inf(), |
516 | 65 | BrickAssertion::budget_met(), |
517 | | ] |
518 | 65 | } |
519 | | } |
520 | | |
521 | | /// QKV projection brick. |
522 | | #[derive(Debug)] |
523 | | pub struct QkvBrick { |
524 | | /// Hidden dimension |
525 | | pub hidden_dim: usize, |
526 | | /// Q output dimension |
527 | | pub q_dim: usize, |
528 | | /// K output dimension |
529 | | pub k_dim: usize, |
530 | | /// V output dimension |
531 | | pub v_dim: usize, |
532 | | /// Budget |
533 | | budget: TokenBudget, |
534 | | /// Has bias (Qwen2 has large biases) |
535 | | pub has_bias: bool, |
536 | | } |
537 | | |
538 | | impl QkvBrick { |
539 | | /// Create a new QKV brick. |
540 | 5 | pub fn new(hidden_dim: usize, q_dim: usize, k_dim: usize, v_dim: usize) -> Self { |
541 | 5 | Self { |
542 | 5 | hidden_dim, |
543 | 5 | q_dim, |
544 | 5 | k_dim, |
545 | 5 | v_dim, |
546 | 5 | budget: TokenBudget::from_latency(6.0), // 6µs target |
547 | 5 | has_bias: false, |
548 | 5 | } |
549 | 5 | } |
550 | | |
551 | | /// Set custom budget. |
552 | | #[must_use] |
553 | 0 | pub fn with_budget(mut self, budget: TokenBudget) -> Self { |
554 | 0 | self.budget = budget; |
555 | 0 | self |
556 | 0 | } |
557 | | |
558 | | /// Mark as having bias. |
559 | | #[must_use] |
560 | 1 | pub fn with_bias(mut self) -> Self { |
561 | 1 | self.has_bias = true; |
562 | 1 | self |
563 | 1 | } |
564 | | |
565 | | /// Total output dimension. |
566 | 1 | pub fn total_out_dim(&self) -> usize { |
567 | 1 | self.q_dim + self.k_dim + self.v_dim |
568 | 1 | } |
569 | | } |
570 | | |
571 | | impl ComputeBrick for QkvBrick { |
572 | | type Output = (Vec<f32>, Vec<f32>, Vec<f32>); |
573 | | |
574 | 0 | fn name(&self) -> &'static str { |
575 | 0 | "qkv_proj" |
576 | 0 | } |
577 | | |
578 | 4 | fn budget(&self) -> TokenBudget { |
579 | 4 | self.budget |
580 | 4 | } |
581 | | |
582 | 0 | fn assertions(&self) -> Vec<BrickAssertion> { |
583 | 0 | vec![ |
584 | 0 | BrickAssertion::no_nan(), |
585 | 0 | BrickAssertion::no_inf(), |
586 | 0 | BrickAssertion::budget_met(), |
587 | | ] |
588 | 0 | } |
589 | | } |
590 | | |
591 | | /// RoPE brick - rotary position embedding. |
592 | | #[derive(Debug)] |
593 | | pub struct RopeBrick { |
594 | | /// Head dimension |
595 | | pub head_dim: usize, |
596 | | /// Number of heads |
597 | | pub num_heads: usize, |
598 | | /// Base theta |
599 | | pub theta: f32, |
600 | | /// RoPE type (0=NORM, 2=NEOX) |
601 | | pub rope_type: u32, |
602 | | /// Budget |
603 | | budget: TokenBudget, |
604 | | } |
605 | | |
606 | | impl RopeBrick { |
607 | | /// Create a new RoPE brick. |
608 | 3 | pub fn new(head_dim: usize, num_heads: usize, theta: f32, rope_type: u32) -> Self { |
609 | 3 | Self { |
610 | 3 | head_dim, |
611 | 3 | num_heads, |
612 | 3 | theta, |
613 | 3 | rope_type, |
614 | 3 | budget: TokenBudget::from_latency(1.0), // 1µs target |
615 | 3 | } |
616 | 3 | } |
617 | | |
618 | | /// Set custom budget. |
619 | | #[must_use] |
620 | 0 | pub fn with_budget(mut self, budget: TokenBudget) -> Self { |
621 | 0 | self.budget = budget; |
622 | 0 | self |
623 | 0 | } |
624 | | } |
625 | | |
626 | | impl ComputeBrick for RopeBrick { |
627 | | type Output = Vec<f32>; |
628 | | |
629 | 0 | fn name(&self) -> &'static str { |
630 | 0 | "rope" |
631 | 0 | } |
632 | | |
633 | 4 | fn budget(&self) -> TokenBudget { |
634 | 4 | self.budget |
635 | 4 | } |
636 | | |
637 | 0 | fn assertions(&self) -> Vec<BrickAssertion> { |
638 | 0 | vec![ |
639 | 0 | BrickAssertion::no_nan(), |
640 | 0 | BrickAssertion::no_inf(), |
641 | 0 | BrickAssertion::budget_met(), |
642 | | ] |
643 | 0 | } |
644 | | } |
645 | | |
646 | | /// Attention brick. |
647 | | #[derive(Debug)] |
648 | | pub struct AttentionBrick { |
649 | | /// Number of query heads |
650 | | pub num_heads: usize, |
651 | | /// Number of KV heads (for GQA) |
652 | | pub num_kv_heads: usize, |
653 | | /// Head dimension |
654 | | pub head_dim: usize, |
655 | | /// Budget |
656 | | budget: TokenBudget, |
657 | | } |
658 | | |
659 | | impl AttentionBrick { |
660 | | /// Create a new attention brick. |
661 | 106 | pub fn new(num_heads: usize, num_kv_heads: usize, head_dim: usize) -> Self { |
662 | 106 | Self { |
663 | 106 | num_heads, |
664 | 106 | num_kv_heads, |
665 | 106 | head_dim, |
666 | 106 | budget: TokenBudget::from_latency(10.0), // 10µs target |
667 | 106 | } |
668 | 106 | } |
669 | | |
670 | | /// Set custom budget. |
671 | | #[must_use] |
672 | 0 | pub fn with_budget(mut self, budget: TokenBudget) -> Self { |
673 | 0 | self.budget = budget; |
674 | 0 | self |
675 | 0 | } |
676 | | |
677 | | /// GQA group size. |
678 | 0 | pub fn group_size(&self) -> usize { |
679 | 0 | self.num_heads / self.num_kv_heads.max(1) |
680 | 0 | } |
681 | | } |
682 | | |
683 | | impl ComputeBrick for AttentionBrick { |
684 | | type Output = Vec<f32>; |
685 | | |
686 | 0 | fn name(&self) -> &'static str { |
687 | 0 | "attention" |
688 | 0 | } |
689 | | |
690 | 7 | fn budget(&self) -> TokenBudget { |
691 | 7 | self.budget |
692 | 7 | } |
693 | | |
694 | 2 | fn assertions(&self) -> Vec<BrickAssertion> { |
695 | 2 | vec![ |
696 | 2 | BrickAssertion::no_nan(), |
697 | 2 | BrickAssertion::no_inf(), |
698 | 2 | BrickAssertion::budget_met(), |
699 | | // Attention outputs should be bounded |
700 | 2 | BrickAssertion::bounds(-100.0, 100.0), |
701 | | ] |
702 | 2 | } |
703 | | } |
704 | | |
705 | | /// FlashAttentionBrick - incremental flash attention for decode (P1 optimization). |
706 | | /// |
707 | | /// **Algorithm** (FlashAttention-2, Dao et al. 2023): |
708 | | /// ```text |
709 | | /// For decode (single query token): |
710 | | /// Q: [1, H, D] (single query) |
711 | | /// K: [S, H_kv, D] (full KV cache) |
712 | | /// V: [S, H_kv, D] (full KV cache) |
713 | | /// |
714 | | /// Online softmax (no full attention matrix materialization): |
715 | | /// for tile in KV_tiles(TILE_SIZE=128): |
716 | | /// S_tile = Q @ K_tile^T / sqrt(D) # [1, H, TILE_SIZE] |
717 | | /// m_new = max(m_old, max(S_tile)) # Running max |
718 | | /// P_tile = exp(S_tile - m_new) # Stable softmax numerator |
719 | | /// O = O * exp(m_old - m_new) + P_tile @ V_tile # Accumulate |
720 | | /// l = l * exp(m_old - m_new) + sum(P_tile) # Running denominator |
721 | | /// O = O / l # Final output |
722 | | /// ``` |
723 | | /// |
724 | | /// **Performance vs naive**: |
725 | | /// - Naive: O(S) memory for attention matrix |
726 | | /// - Flash: O(TILE_SIZE) memory, 2x speedup from better cache locality |
727 | | /// |
728 | | /// **Reference**: Dao, T., et al. (2023). "FlashAttention-2: Faster Attention |
729 | | /// with Better Parallelism and Work Partitioning." arXiv:2307.08691. |
730 | | #[derive(Debug, Clone)] |
731 | | pub struct FlashAttentionBrick { |
732 | | /// Number of query heads |
733 | | pub num_heads: usize, |
734 | | /// Number of KV heads (for GQA) |
735 | | pub num_kv_heads: usize, |
736 | | /// Head dimension |
737 | | pub head_dim: usize, |
738 | | /// Tile size for KV cache (default: 128 for L2 cache fit) |
739 | | pub tile_size: usize, |
740 | | /// Budget (target: 5.0µs for 2x improvement over naive) |
741 | | budget: TokenBudget, |
742 | | /// Use online softmax (FlashAttention algorithm) |
743 | | pub use_online_softmax: bool, |
744 | | } |
745 | | |
746 | | impl FlashAttentionBrick { |
747 | | /// Create new flash attention brick with default tile size. |
748 | | #[must_use] |
749 | 15 | pub fn new(num_heads: usize, num_kv_heads: usize, head_dim: usize) -> Self { |
750 | 15 | Self { |
751 | 15 | num_heads, |
752 | 15 | num_kv_heads, |
753 | 15 | head_dim, |
754 | 15 | tile_size: 128, // Optimal for L2 cache on most GPUs |
755 | 15 | budget: TokenBudget::from_latency(5.0), // 5µs target (2x vs 10µs naive) |
756 | 15 | use_online_softmax: true, |
757 | 15 | } |
758 | 15 | } |
759 | | |
760 | | /// Create with custom tile size (for tuning). |
761 | | #[must_use] |
762 | 1 | pub fn with_tile_size( |
763 | 1 | num_heads: usize, |
764 | 1 | num_kv_heads: usize, |
765 | 1 | head_dim: usize, |
766 | 1 | tile_size: usize, |
767 | 1 | ) -> Self { |
768 | 1 | Self { |
769 | 1 | num_heads, |
770 | 1 | num_kv_heads, |
771 | 1 | head_dim, |
772 | 1 | tile_size, |
773 | 1 | budget: TokenBudget::from_latency(5.0), |
774 | 1 | use_online_softmax: true, |
775 | 1 | } |
776 | 1 | } |
777 | | |
778 | | /// Set custom budget. |
779 | | #[must_use] |
780 | 0 | pub fn with_budget(mut self, budget: TokenBudget) -> Self { |
781 | 0 | self.budget = budget; |
782 | 0 | self |
783 | 0 | } |
784 | | |
785 | | /// GQA group size (query heads per KV head). |
786 | | #[must_use] |
787 | 5 | pub fn group_size(&self) -> usize { |
788 | 5 | self.num_heads / self.num_kv_heads.max(1) |
789 | 5 | } |
790 | | |
791 | | /// Compute FLOPs per token for attention at given sequence length. |
792 | | /// |
793 | | /// FLOPs = 2 * H * D * S (Q @ K^T) + 2 * H * S * D (attn @ V) |
794 | | /// = 4 * H * D * S |
795 | | #[must_use] |
796 | 2 | pub fn flops(&self, seq_len: usize) -> u64 { |
797 | 2 | 4 * self.num_heads as u64 * self.head_dim as u64 * seq_len as u64 |
798 | 2 | } |
799 | | |
800 | | /// Compute memory bandwidth for naive vs flash attention. |
801 | | /// |
802 | | /// Naive: Reads full KV cache for each head |
803 | | /// Flash: Reads KV cache in tiles, better cache reuse |
804 | | #[must_use] |
805 | 4 | pub fn memory_bytes(&self, seq_len: usize) -> (u64, u64) { |
806 | 4 | let kv_bytes = 2 * self.num_kv_heads as u64 * self.head_dim as u64 * seq_len as u64 * 4; // K + V, f32 |
807 | 4 | let naive = kv_bytes + self.num_heads as u64 * seq_len as u64 * 4; // + attention matrix |
808 | 4 | let flash = kv_bytes; // No attention matrix materialized |
809 | 4 | (naive, flash) |
810 | 4 | } |
811 | | |
812 | | /// Compute arithmetic intensity (FLOPs / bytes). |
813 | | #[must_use] |
814 | 0 | pub fn arithmetic_intensity(&self, seq_len: usize) -> f64 { |
815 | 0 | let (_, flash_bytes) = self.memory_bytes(seq_len); |
816 | 0 | self.flops(seq_len) as f64 / flash_bytes as f64 |
817 | 0 | } |
818 | | |
819 | | /// Number of tiles needed for given sequence length. |
820 | | #[must_use] |
821 | 3 | pub fn num_tiles(&self, seq_len: usize) -> usize { |
822 | 3 | seq_len.div_ceil(self.tile_size) |
823 | 3 | } |
824 | | |
825 | | /// Compute flash attention with online softmax algorithm. |
826 | | /// |
827 | | /// **REAL IMPLEMENTATION** - FlashAttention-2 (Dao et al. 2023) |
828 | | /// |
829 | | /// # Arguments |
830 | | /// * `query` - Query tensor [num_heads, head_dim] |
831 | | /// * `keys` - Key cache [seq_len, num_kv_heads, head_dim] |
832 | | /// * `values` - Value cache [seq_len, num_kv_heads, head_dim] |
833 | | /// |
834 | | /// # Returns |
835 | | /// * Output tensor [num_heads, head_dim] |
836 | 4 | pub fn forward( |
837 | 4 | &self, |
838 | 4 | query: &[f32], // [num_heads * head_dim] |
839 | 4 | keys: &[f32], // [seq_len * num_kv_heads * head_dim] |
840 | 4 | values: &[f32], // [seq_len * num_kv_heads * head_dim] |
841 | 4 | seq_len: usize, |
842 | 4 | ) -> Result<Vec<f32>, BrickError> { |
843 | 4 | if self.num_heads == 0 || self.head_dim == 0 { |
844 | 0 | return Err(BrickError::InvalidInput("Zero dimension".to_string())); |
845 | 4 | } |
846 | 4 | if query.len() != self.num_heads * self.head_dim { |
847 | 0 | return Err(BrickError::InvalidInput(format!( |
848 | 0 | "Query length {} != num_heads * head_dim = {}", |
849 | 0 | query.len(), |
850 | 0 | self.num_heads * self.head_dim |
851 | 0 | ))); |
852 | 4 | } |
853 | 4 | let expected_kv_len = seq_len * self.num_kv_heads * self.head_dim; |
854 | 4 | if keys.len() != expected_kv_len || values.len() != expected_kv_len { |
855 | 0 | return Err(BrickError::InvalidInput(format!( |
856 | 0 | "KV length {} != seq_len * num_kv_heads * head_dim = {}", |
857 | 0 | keys.len(), |
858 | 0 | expected_kv_len |
859 | 0 | ))); |
860 | 4 | } |
861 | | |
862 | 4 | let scale = 1.0 / (self.head_dim as f32).sqrt(); |
863 | 4 | let group_size = self.group_size(); |
864 | 4 | let mut output = vec![0.0f32; self.num_heads * self.head_dim]; |
865 | | |
866 | | // Process each query head |
867 | 14 | for h in 0..self.num_heads4 { |
868 | 14 | let kv_head = h / group_size; // GQA: map query head to KV head |
869 | 14 | let q_start = h * self.head_dim; |
870 | | |
871 | | // Online softmax variables (FlashAttention-2) |
872 | 14 | let mut m = f32::NEG_INFINITY; // Running max |
873 | 14 | let mut l = 0.0f32; // Running sum of exp |
874 | 14 | let mut o = vec![0.0f32; self.head_dim]; // Running output |
875 | | |
876 | | // Process in tiles for cache efficiency |
877 | 14 | for tile_start in (0..seq_len).step_by(self.tile_size) { |
878 | 14 | let tile_end = (tile_start + self.tile_size).min(seq_len); |
879 | | |
880 | 1.04k | for s in tile_start14 ..tile_end14 { |
881 | | // Compute Q @ K^T for this position |
882 | 1.04k | let k_start = (s * self.num_kv_heads + kv_head) * self.head_dim; |
883 | 1.04k | let mut score = 0.0f32; |
884 | 65.6k | for d in 0..self.head_dim1.04k { |
885 | 65.6k | score += query[q_start + d] * keys[k_start + d]; |
886 | 65.6k | } |
887 | 1.04k | score *= scale; |
888 | | |
889 | | // Online softmax update (Milakov & Gimelshein, 2018) |
890 | 1.04k | let m_new = m.max(score); |
891 | 1.04k | let exp_old = (m - m_new).exp(); |
892 | 1.04k | let exp_score = (score - m_new).exp(); |
893 | | |
894 | | // Update running sum |
895 | 1.04k | l = l * exp_old + exp_score; |
896 | | |
897 | | // Update running output: O = O * exp(m_old - m_new) + exp(score - m_new) * V |
898 | 1.04k | let v_start = (s * self.num_kv_heads + kv_head) * self.head_dim; |
899 | 65.6k | for d in 0..self.head_dim1.04k { |
900 | 65.6k | o[d] = o[d] * exp_old + exp_score * values[v_start + d]; |
901 | 65.6k | } |
902 | | |
903 | 1.04k | m = m_new; |
904 | | } |
905 | | } |
906 | | |
907 | | // Normalize output: O = O / l |
908 | 14 | if l > 0.0 { |
909 | 552 | for d in 0..self.head_dim14 { |
910 | 552 | output[h * self.head_dim + d] = o[d] / l; |
911 | 552 | } |
912 | 0 | } |
913 | | } |
914 | | |
915 | 4 | Ok(output) |
916 | 4 | } |
917 | | |
918 | | /// Execute flash attention with timing (for benchmarking). |
919 | 1 | pub fn forward_timed( |
920 | 1 | &self, |
921 | 1 | query: &[f32], |
922 | 1 | keys: &[f32], |
923 | 1 | values: &[f32], |
924 | 1 | seq_len: usize, |
925 | 1 | ) -> Result<TokenResult<Vec<f32>>, BrickError> { |
926 | 1 | let start = Instant::now(); |
927 | 1 | let output = self.forward(query, keys, values, seq_len)?0 ; |
928 | 1 | let elapsed_us = start.elapsed().as_secs_f64() * 1_000_000.0; |
929 | | |
930 | 1 | Ok(TokenResult { |
931 | 1 | output, |
932 | 1 | tokens_processed: 1, |
933 | 1 | us_per_token: elapsed_us, |
934 | 1 | tokens_per_sec: 1_000_000.0 / elapsed_us, |
935 | 1 | budget_met: elapsed_us <= self.budget.us_per_token, |
936 | 1 | }) |
937 | 1 | } |
938 | | |
939 | | /// Legacy stub for backward compatibility (prefer `forward()`) |
940 | | #[deprecated(note = "Use forward() for real implementation")] |
941 | 1 | pub fn execute(&self, _seq_len: usize) -> Result<Vec<f32>, BrickError> { |
942 | 1 | if self.num_heads == 0 || self.head_dim == 00 { |
943 | 1 | return Err(BrickError::InvalidInput("Zero dimension".to_string())); |
944 | 0 | } |
945 | 0 | Ok(vec![0.0; self.num_heads * self.head_dim]) |
946 | 1 | } |
947 | | } |
948 | | |
949 | | impl ComputeBrick for FlashAttentionBrick { |
950 | | type Output = Vec<f32>; |
951 | | |
952 | 0 | fn name(&self) -> &'static str { |
953 | 0 | "flash_attention" |
954 | 0 | } |
955 | | |
956 | 1 | fn budget(&self) -> TokenBudget { |
957 | 1 | self.budget |
958 | 1 | } |
959 | | |
960 | 1 | fn assertions(&self) -> Vec<BrickAssertion> { |
961 | 1 | vec![ |
962 | 1 | BrickAssertion::no_nan(), |
963 | 1 | BrickAssertion::no_inf(), |
964 | 1 | BrickAssertion::budget_met(), |
965 | | // Attention outputs should be bounded |
966 | 1 | BrickAssertion::bounds(-100.0, 100.0), |
967 | | // Custom assertions for flash attention |
968 | 1 | BrickAssertion { |
969 | 1 | name: "online_softmax".to_string(), |
970 | 1 | description: "Uses online softmax (no full attention matrix)".to_string(), |
971 | 1 | kind: AssertionKind::Custom { |
972 | 1 | check_name: "online_softmax".to_string(), |
973 | 1 | }, |
974 | 1 | }, |
975 | 1 | BrickAssertion { |
976 | 1 | name: "tiled_kv_access".to_string(), |
977 | 1 | description: "KV cache accessed in tiles for cache locality".to_string(), |
978 | 1 | kind: AssertionKind::Custom { |
979 | 1 | check_name: "tiled_kv_access".to_string(), |
980 | 1 | }, |
981 | 1 | }, |
982 | | ] |
983 | 1 | } |
984 | | |
985 | 0 | fn can_run(&self) -> bool { |
986 | 0 | self.num_heads > 0 && self.head_dim > 0 && self.tile_size > 0 |
987 | 0 | } |
988 | | } |
989 | | |
990 | | /// FFN brick (SwiGLU). |
991 | | #[derive(Debug)] |
992 | | pub struct FfnBrick { |
993 | | /// Hidden dimension |
994 | | pub hidden_dim: usize, |
995 | | /// Intermediate dimension |
996 | | pub intermediate_dim: usize, |
997 | | /// Budget |
998 | | budget: TokenBudget, |
999 | | } |
1000 | | |
1001 | | impl FfnBrick { |
1002 | | /// Create a new FFN brick. |
1003 | 105 | pub fn new(hidden_dim: usize, intermediate_dim: usize) -> Self { |
1004 | 105 | Self { |
1005 | 105 | hidden_dim, |
1006 | 105 | intermediate_dim, |
1007 | 105 | budget: TokenBudget::from_latency(12.2), // 12.2µs target |
1008 | 105 | } |
1009 | 105 | } |
1010 | | |
1011 | | /// Set custom budget. |
1012 | | #[must_use] |
1013 | 0 | pub fn with_budget(mut self, budget: TokenBudget) -> Self { |
1014 | 0 | self.budget = budget; |
1015 | 0 | self |
1016 | 0 | } |
1017 | | } |
1018 | | |
1019 | | impl ComputeBrick for FfnBrick { |
1020 | | type Output = Vec<f32>; |
1021 | | |
1022 | 0 | fn name(&self) -> &'static str { |
1023 | 0 | "ffn" |
1024 | 0 | } |
1025 | | |
1026 | 6 | fn budget(&self) -> TokenBudget { |
1027 | 6 | self.budget |
1028 | 6 | } |
1029 | | |
1030 | 2 | fn assertions(&self) -> Vec<BrickAssertion> { |
1031 | 2 | vec![ |
1032 | 2 | BrickAssertion::no_nan(), |
1033 | 2 | BrickAssertion::no_inf(), |
1034 | 2 | BrickAssertion::budget_met(), |
1035 | | ] |
1036 | 2 | } |
1037 | | } |
1038 | | |
1039 | | /// Output projection brick. |
1040 | | #[derive(Debug)] |
1041 | | pub struct OProjBrick { |
1042 | | /// Input dimension (num_heads * head_dim) |
1043 | | pub in_dim: usize, |
1044 | | /// Output dimension (hidden_dim) |
1045 | | pub out_dim: usize, |
1046 | | /// Budget |
1047 | | budget: TokenBudget, |
1048 | | } |
1049 | | |
1050 | | impl OProjBrick { |
1051 | | /// Create a new O projection brick. |
1052 | 2 | pub fn new(in_dim: usize, out_dim: usize) -> Self { |
1053 | 2 | Self { |
1054 | 2 | in_dim, |
1055 | 2 | out_dim, |
1056 | 2 | budget: TokenBudget::from_latency(3.5), // 3.5µs target |
1057 | 2 | } |
1058 | 2 | } |
1059 | | |
1060 | | /// Set custom budget. |
1061 | | #[must_use] |
1062 | 0 | pub fn with_budget(mut self, budget: TokenBudget) -> Self { |
1063 | 0 | self.budget = budget; |
1064 | 0 | self |
1065 | 0 | } |
1066 | | } |
1067 | | |
1068 | | impl ComputeBrick for OProjBrick { |
1069 | | type Output = Vec<f32>; |
1070 | | |
1071 | 0 | fn name(&self) -> &'static str { |
1072 | 0 | "o_proj" |
1073 | 0 | } |
1074 | | |
1075 | 3 | fn budget(&self) -> TokenBudget { |
1076 | 3 | self.budget |
1077 | 3 | } |
1078 | | |
1079 | 0 | fn assertions(&self) -> Vec<BrickAssertion> { |
1080 | 0 | vec![ |
1081 | 0 | BrickAssertion::no_nan(), |
1082 | 0 | BrickAssertion::no_inf(), |
1083 | 0 | BrickAssertion::budget_met(), |
1084 | | ] |
1085 | 0 | } |
1086 | | } |
1087 | | |
1088 | | /// ActivationQuantBrick - Q8 activation quantization for memory bandwidth reduction (P2). |
1089 | | /// |
1090 | | /// **Purpose**: Quantize intermediate activations from f32 to int8 to reduce |
1091 | | /// memory bandwidth by 4x during inter-layer communication. |
1092 | | /// |
1093 | | /// **Pipeline**: |
1094 | | /// ```text |
1095 | | /// Layer N output (f32) → Q8 quantize → transfer → Q8 dequantize → Layer N+1 input (f32) |
1096 | | /// |
1097 | | /// Memory bandwidth reduction: |
1098 | | /// f32: 4 bytes/element |
1099 | | /// int8: 1 byte/element + 2 floats (scale, zero_point) |
1100 | | /// Effective: ~4x reduction for large activations |
1101 | | /// ``` |
1102 | | /// |
1103 | | /// **Algorithm** (per-tensor affine quantization): |
1104 | | /// ```text |
1105 | | /// Quantize: |
1106 | | /// scale = (max - min) / 255 |
1107 | | /// zero_point = round(-min / scale) |
1108 | | /// q[i] = clamp(round(x[i] / scale + zero_point), 0, 255) |
1109 | | /// |
1110 | | /// Dequantize: |
1111 | | /// x[i] = (q[i] - zero_point) * scale |
1112 | | /// ``` |
1113 | | /// |
1114 | | /// **Performance**: 2x memory BW improvement with ~0.1% accuracy loss |
1115 | | /// |
1116 | | /// **Reference**: Jacob, B., et al. (2018). "Quantization and Training of |
1117 | | /// Neural Networks for Efficient Integer-Arithmetic-Only Inference." CVPR '18. |
1118 | | #[derive(Debug, Clone)] |
1119 | | pub struct ActivationQuantBrick { |
1120 | | /// Activation dimension (e.g., hidden_dim or intermediate_dim) |
1121 | | pub dim: usize, |
1122 | | /// Budget (target: 0.5µs for quant+dequant overhead) |
1123 | | budget: TokenBudget, |
1124 | | /// Use per-channel quantization (more accurate but slower) |
1125 | | pub per_channel: bool, |
1126 | | } |
1127 | | |
1128 | | impl ActivationQuantBrick { |
1129 | | /// Create new activation quantization brick. |
1130 | | #[must_use] |
1131 | 14 | pub fn new(dim: usize) -> Self { |
1132 | 14 | Self { |
1133 | 14 | dim, |
1134 | 14 | budget: TokenBudget::from_latency(0.5), // 0.5µs overhead target |
1135 | 14 | per_channel: false, |
1136 | 14 | } |
1137 | 14 | } |
1138 | | |
1139 | | /// Create with per-channel quantization. |
1140 | | #[must_use] |
1141 | 1 | pub fn with_per_channel(dim: usize) -> Self { |
1142 | 1 | Self { |
1143 | 1 | dim, |
1144 | 1 | budget: TokenBudget::from_latency(1.0), // 1.0µs for per-channel |
1145 | 1 | per_channel: true, |
1146 | 1 | } |
1147 | 1 | } |
1148 | | |
1149 | | /// Set custom budget. |
1150 | | #[must_use] |
1151 | 0 | pub fn with_budget(mut self, budget: TokenBudget) -> Self { |
1152 | 0 | self.budget = budget; |
1153 | 0 | self |
1154 | 0 | } |
1155 | | |
1156 | | /// Compute memory bandwidth reduction factor. |
1157 | | /// |
1158 | | /// f32 (4 bytes) → int8 (1 byte) + scale/zero_point = ~4x reduction |
1159 | | #[must_use] |
1160 | 1 | pub fn bandwidth_reduction(&self) -> f64 { |
1161 | | // Original: dim * 4 bytes (f32) |
1162 | | // Quantized: dim * 1 byte (int8) + 8 bytes (scale + zero_point) |
1163 | 1 | let original_bytes = self.dim * 4; |
1164 | 1 | let quantized_bytes = self.dim + 8; // +8 for scale and zero_point (f32 each) |
1165 | 1 | original_bytes as f64 / quantized_bytes as f64 |
1166 | 1 | } |
1167 | | |
1168 | | /// Compute quantization error estimate (typical for 8-bit). |
1169 | | /// |
1170 | | /// Per Jacob et al. 2018, typical Q8 error is ~0.1% for activations. |
1171 | | #[must_use] |
1172 | 2 | pub fn estimated_error(&self) -> f64 { |
1173 | 2 | if self.per_channel { |
1174 | 1 | 0.0005 // 0.05% for per-channel |
1175 | | } else { |
1176 | 1 | 0.001 // 0.1% for per-tensor |
1177 | | } |
1178 | 2 | } |
1179 | | |
1180 | | /// Compute bytes saved per token. |
1181 | | #[must_use] |
1182 | 3 | pub fn bytes_saved(&self) -> usize { |
1183 | | // f32 (4 bytes) → int8 (1 byte) = 3 bytes saved per element |
1184 | 3 | self.dim * 3 |
1185 | 3 | } |
1186 | | |
1187 | | /// Quantize f32 activations to int8 using Q8_0 block format. |
1188 | | /// |
1189 | | /// **REAL IMPLEMENTATION** - Not a stub. |
1190 | | /// Uses symmetric quantization: scale = max(abs(values)) / 127.0 |
1191 | | /// |
1192 | | /// # Arguments |
1193 | | /// * `input` - f32 activations to quantize (must be length == self.dim) |
1194 | | /// |
1195 | | /// # Returns |
1196 | | /// * Quantized int8 values and scale factors |
1197 | | /// |
1198 | | /// # Example |
1199 | | /// ```ignore |
1200 | | /// let brick = ActivationQuantBrick::new(64); |
1201 | | /// let input = vec![1.0f32; 64]; |
1202 | | /// let (quants, scales) = brick.quantize(&input)?; |
1203 | | /// assert_eq!(quants.len(), 64); |
1204 | | /// ``` |
1205 | 4 | pub fn quantize(&self, input: &[f32]) -> Result<(Vec<i8>, Vec<f32>), BrickError> { |
1206 | 4 | if input.len() != self.dim { |
1207 | 1 | return Err(BrickError::InvalidInput(format!( |
1208 | 1 | "Input length {} != dim {}", |
1209 | 1 | input.len(), |
1210 | 1 | self.dim |
1211 | 1 | ))); |
1212 | 3 | } |
1213 | 3 | if self.dim == 0 { |
1214 | 0 | return Err(BrickError::InvalidInput("Zero dimension".to_string())); |
1215 | 3 | } |
1216 | | |
1217 | | // Quantize in blocks of 32 (Q8_0 block size) |
1218 | 3 | let num_blocks = self.dim.div_ceil(32); |
1219 | 3 | let mut quants = Vec::with_capacity(self.dim); |
1220 | 3 | let mut scales = Vec::with_capacity(num_blocks); |
1221 | | |
1222 | 35 | for block_idx in 0..num_blocks3 { |
1223 | 35 | let start = block_idx * 32; |
1224 | 35 | let end = (start + 32).min(self.dim); |
1225 | | |
1226 | | // Pad to 32 if needed |
1227 | 35 | let mut block_data = [0.0f32; 32]; |
1228 | 1.12k | for (i, &v) in input35 [start..end].iter35 ().enumerate35 () { |
1229 | 1.12k | block_data[i] = v; |
1230 | 1.12k | } |
1231 | | |
1232 | 35 | let block = Q8_0Block::quantize(&block_data); |
1233 | 35 | scales.push(block.scale); |
1234 | | |
1235 | | // Only take the actual values (not padding) |
1236 | 1.12k | for &q in &block.quants35 [0..(end - start)]35 { |
1237 | 1.12k | quants.push(q); |
1238 | 1.12k | } |
1239 | | } |
1240 | | |
1241 | 3 | Ok((quants, scales)) |
1242 | 4 | } |
1243 | | |
1244 | | /// Dequantize int8 back to f32 using stored scales. |
1245 | | /// |
1246 | | /// **REAL IMPLEMENTATION** - Not a stub. |
1247 | 2 | pub fn dequantize(&self, quants: &[i8], scales: &[f32]) -> Result<Vec<f32>, BrickError> { |
1248 | 2 | if quants.len() != self.dim { |
1249 | 0 | return Err(BrickError::InvalidInput(format!( |
1250 | 0 | "Quants length {} != dim {}", |
1251 | 0 | quants.len(), |
1252 | 0 | self.dim |
1253 | 0 | ))); |
1254 | 2 | } |
1255 | | |
1256 | 2 | let mut output = Vec::with_capacity(self.dim); |
1257 | 2 | for (block_idx, &scale) in scales.iter().enumerate() { |
1258 | 2 | let start = block_idx * 32; |
1259 | 2 | let end = (start + 32).min(self.dim); |
1260 | 64 | for &q in &quants2 [start..end]2 { |
1261 | 64 | output.push(q as f32 * scale); |
1262 | 64 | } |
1263 | | } |
1264 | | |
1265 | 2 | Ok(output) |
1266 | 2 | } |
1267 | | |
1268 | | /// Compute quantization error vs original input. |
1269 | | /// |
1270 | | /// **REAL IMPLEMENTATION** - Measures actual error, not estimates. |
1271 | 1 | pub fn measure_error( |
1272 | 1 | &self, |
1273 | 1 | original: &[f32], |
1274 | 1 | quants: &[i8], |
1275 | 1 | scales: &[f32], |
1276 | 1 | ) -> Result<f64, BrickError> { |
1277 | 1 | let dequantized = self.dequantize(quants, scales)?0 ; |
1278 | | |
1279 | 1 | let max_error = original |
1280 | 1 | .iter() |
1281 | 1 | .zip(dequantized.iter()) |
1282 | 32 | .map1 (|(a, b)| (a - b).abs()) |
1283 | 1 | .fold(0.0f32, f32::max); |
1284 | | |
1285 | 32 | let max_val1 = original1 .iter1 ().map1 (|v| v.abs()).fold1 (0.0f32, f32::max); |
1286 | 1 | if max_val < 1e-10 { |
1287 | 0 | return Ok(0.0); |
1288 | 1 | } |
1289 | | |
1290 | 1 | Ok((max_error / max_val) as f64) |
1291 | 1 | } |
1292 | | |
1293 | | /// Execute quantization with timing (for benchmarking). |
1294 | | #[allow(clippy::type_complexity)] |
1295 | 1 | pub fn execute_timed( |
1296 | 1 | &self, |
1297 | 1 | input: &[f32], |
1298 | 1 | ) -> Result<TokenResult<(Vec<i8>, Vec<f32>)>, BrickError> { |
1299 | 1 | let start = Instant::now(); |
1300 | 1 | let (quants, scales) = self.quantize(input)?0 ; |
1301 | 1 | let elapsed_us = start.elapsed().as_secs_f64() * 1_000_000.0; |
1302 | | |
1303 | 1 | Ok(TokenResult { |
1304 | 1 | output: (quants, scales), |
1305 | 1 | tokens_processed: 1, |
1306 | 1 | us_per_token: elapsed_us, |
1307 | 1 | tokens_per_sec: 1_000_000.0 / elapsed_us, |
1308 | 1 | budget_met: elapsed_us <= self.budget.us_per_token, |
1309 | 1 | }) |
1310 | 1 | } |
1311 | | |
1312 | | /// Legacy stub for backward compatibility (prefer `quantize()`) |
1313 | | #[deprecated(note = "Use quantize() for real implementation")] |
1314 | 2 | pub fn execute(&self) -> Result<Vec<u8>, BrickError> { |
1315 | 2 | if self.dim == 0 { |
1316 | 2 | return Err(BrickError::InvalidInput("Zero dimension".to_string())); |
1317 | 0 | } |
1318 | | // Return zeros for backward compat - use quantize() for real output |
1319 | 0 | Ok(vec![128u8; self.dim]) |
1320 | 2 | } |
1321 | | } |
1322 | | |
1323 | | impl ComputeBrick for ActivationQuantBrick { |
1324 | | type Output = Vec<u8>; |
1325 | | |
1326 | 1 | fn name(&self) -> &'static str { |
1327 | 1 | "activation_quant" |
1328 | 1 | } |
1329 | | |
1330 | 1 | fn budget(&self) -> TokenBudget { |
1331 | 1 | self.budget |
1332 | 1 | } |
1333 | | |
1334 | 1 | fn assertions(&self) -> Vec<BrickAssertion> { |
1335 | 1 | vec![ |
1336 | 1 | BrickAssertion::budget_met(), |
1337 | 1 | BrickAssertion { |
1338 | 1 | name: "symmetric_range".to_string(), |
1339 | 1 | description: "Q8 values centered around 128 (zero_point)".to_string(), |
1340 | 1 | kind: AssertionKind::Custom { |
1341 | 1 | check_name: "symmetric_range".to_string(), |
1342 | 1 | }, |
1343 | 1 | }, |
1344 | 1 | BrickAssertion { |
1345 | 1 | name: "error_bound".to_string(), |
1346 | 1 | description: "Quantization error < 0.1% (per-tensor) or 0.05% (per-channel)" |
1347 | 1 | .to_string(), |
1348 | 1 | kind: AssertionKind::Custom { |
1349 | 1 | check_name: "error_bound".to_string(), |
1350 | 1 | }, |
1351 | 1 | }, |
1352 | | ] |
1353 | 1 | } |
1354 | | |
1355 | 2 | fn can_run(&self) -> bool { |
1356 | 2 | self.dim > 0 |
1357 | 2 | } |
1358 | | } |
1359 | | |
1360 | | // ============================================================================ |
1361 | | // Transformer Layer Brick |
1362 | | // ============================================================================ |
1363 | | |
1364 | | /// Full transformer layer as a composed brick. |
1365 | | #[derive(Debug)] |
1366 | | pub struct TransformerLayerBrick { |
1367 | | /// Layer index |
1368 | | pub layer_idx: usize, |
1369 | | /// Attention layer normalization brick |
1370 | | pub attn_norm: RmsNormBrick, |
1371 | | /// QKV projection brick |
1372 | | pub qkv: QkvBrick, |
1373 | | /// Rotary position embedding brick |
1374 | | pub rope: RopeBrick, |
1375 | | /// Attention computation brick |
1376 | | pub attention: AttentionBrick, |
1377 | | /// Output projection brick |
1378 | | pub o_proj: OProjBrick, |
1379 | | /// FFN layer normalization brick |
1380 | | pub ffn_norm: RmsNormBrick, |
1381 | | /// Feed-forward network brick |
1382 | | pub ffn: FfnBrick, |
1383 | | /// Timing metrics (updated after each run) |
1384 | | pub last_timing: Option<LayerTiming>, |
1385 | | } |
1386 | | |
1387 | | /// Timing breakdown for a layer. |
1388 | | #[derive(Debug, Clone, Default)] |
1389 | | pub struct LayerTiming { |
1390 | | /// Attention normalization time (µs) |
1391 | | pub attn_norm_us: f64, |
1392 | | /// QKV projection time (µs) |
1393 | | pub qkv_us: f64, |
1394 | | /// RoPE application time (µs) |
1395 | | pub rope_us: f64, |
1396 | | /// Attention computation time (µs) |
1397 | | pub attention_us: f64, |
1398 | | /// Output projection time (µs) |
1399 | | pub o_proj_us: f64, |
1400 | | /// FFN normalization time (µs) |
1401 | | pub ffn_norm_us: f64, |
1402 | | /// FFN computation time (µs) |
1403 | | pub ffn_us: f64, |
1404 | | /// Total layer time (µs) |
1405 | | pub total_us: f64, |
1406 | | } |
1407 | | |
1408 | | impl LayerTiming { |
1409 | | /// Find the bottleneck brick. |
1410 | 2 | pub fn bottleneck(&self) -> (&'static str, f64) { |
1411 | 2 | let bricks = [ |
1412 | 2 | ("attn_norm", self.attn_norm_us), |
1413 | 2 | ("qkv", self.qkv_us), |
1414 | 2 | ("rope", self.rope_us), |
1415 | 2 | ("attention", self.attention_us), |
1416 | 2 | ("o_proj", self.o_proj_us), |
1417 | 2 | ("ffn_norm", self.ffn_norm_us), |
1418 | 2 | ("ffn", self.ffn_us), |
1419 | 2 | ]; |
1420 | | |
1421 | 2 | bricks |
1422 | 2 | .into_iter() |
1423 | 12 | .max_by2 (|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal)) |
1424 | 2 | .unwrap_or(("unknown", 0.0)) |
1425 | 2 | } |
1426 | | } |
1427 | | |
1428 | | impl TransformerLayerBrick { |
1429 | | /// Create from configuration. |
1430 | 2 | pub fn from_config( |
1431 | 2 | layer_idx: usize, |
1432 | 2 | hidden_dim: usize, |
1433 | 2 | num_heads: usize, |
1434 | 2 | num_kv_heads: usize, |
1435 | 2 | intermediate_dim: usize, |
1436 | 2 | eps: f32, |
1437 | 2 | rope_theta: f32, |
1438 | 2 | rope_type: u32, |
1439 | 2 | ) -> Self { |
1440 | 2 | let head_dim = hidden_dim / num_heads; |
1441 | 2 | let q_dim = num_heads * head_dim; |
1442 | 2 | let kv_dim = num_kv_heads * head_dim; |
1443 | | |
1444 | 2 | Self { |
1445 | 2 | layer_idx, |
1446 | 2 | attn_norm: RmsNormBrick::new(vec![1.0; hidden_dim], eps), |
1447 | 2 | qkv: QkvBrick::new(hidden_dim, q_dim, kv_dim, kv_dim), |
1448 | 2 | rope: RopeBrick::new(head_dim, num_heads, rope_theta, rope_type), |
1449 | 2 | attention: AttentionBrick::new(num_heads, num_kv_heads, head_dim), |
1450 | 2 | o_proj: OProjBrick::new(q_dim, hidden_dim), |
1451 | 2 | ffn_norm: RmsNormBrick::new(vec![1.0; hidden_dim], eps), |
1452 | 2 | ffn: FfnBrick::new(hidden_dim, intermediate_dim), |
1453 | 2 | last_timing: None, |
1454 | 2 | } |
1455 | 2 | } |
1456 | | |
1457 | | /// Get total budget for this layer. |
1458 | 1 | pub fn total_budget_us(&self) -> f64 { |
1459 | 1 | self.attn_norm.budget().us_per_token |
1460 | 1 | + self.qkv.budget().us_per_token |
1461 | 1 | + self.rope.budget().us_per_token |
1462 | 1 | + self.attention.budget().us_per_token |
1463 | 1 | + self.o_proj.budget().us_per_token |
1464 | 1 | + self.ffn_norm.budget().us_per_token |
1465 | 1 | + self.ffn.budget().us_per_token |
1466 | 1 | } |
1467 | | } |
1468 | | |
1469 | | impl ComputeBrick for TransformerLayerBrick { |
1470 | | type Output = Vec<f32>; |
1471 | | |
1472 | 0 | fn name(&self) -> &'static str { |
1473 | 0 | "transformer_layer" |
1474 | 0 | } |
1475 | | |
1476 | 0 | fn budget(&self) -> TokenBudget { |
1477 | 0 | TokenBudget::from_latency(self.total_budget_us()) |
1478 | 0 | } |
1479 | | |
1480 | 1 | fn assertions(&self) -> Vec<BrickAssertion> { |
1481 | 1 | vec![ |
1482 | 1 | BrickAssertion::no_nan(), |
1483 | 1 | BrickAssertion::no_inf(), |
1484 | 1 | BrickAssertion::budget_met(), |
1485 | | ] |
1486 | 1 | } |
1487 | | |
1488 | 1 | fn verify(&self) -> BrickVerification { |
1489 | | // Verify all component bricks |
1490 | 1 | let mut result = BrickVerification::pass(); |
1491 | | |
1492 | 2 | for brick in [ |
1493 | 1 | &self.attn_norm as &dyn ComputeBrick<Output = Vec<f32>>, |
1494 | 1 | &self.ffn_norm as &dyn ComputeBrick<Output = Vec<f32>>, |
1495 | | ] { |
1496 | 2 | let v = brick.verify(); |
1497 | 2 | if !v.is_valid { |
1498 | 0 | result.is_valid = false; |
1499 | 0 | result.results.extend(v.results); |
1500 | 2 | } |
1501 | | } |
1502 | | |
1503 | 1 | result |
1504 | 1 | } |
1505 | | } |
1506 | | |
1507 | | // ============================================================================ |
1508 | | // Bottleneck Report |
1509 | | // ============================================================================ |
1510 | | |
1511 | | /// Report identifying pipeline bottleneck. |
1512 | | #[derive(Debug, Clone)] |
1513 | | pub struct BottleneckReport { |
1514 | | /// Layer index containing bottleneck |
1515 | | pub layer_idx: usize, |
1516 | | /// Brick name |
1517 | | pub brick_name: &'static str, |
1518 | | /// Actual latency (µs) |
1519 | | pub actual_us: f64, |
1520 | | /// Budget latency (µs) |
1521 | | pub budget_us: f64, |
1522 | | /// Gap factor (actual / budget) |
1523 | | pub gap_factor: f64, |
1524 | | } |
1525 | | |
1526 | | impl fmt::Display for BottleneckReport { |
1527 | 0 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
1528 | 0 | write!( |
1529 | 0 | f, |
1530 | 0 | "Bottleneck: {} (layer {}) - {:.1}µs actual vs {:.1}µs budget ({:.2}x)", |
1531 | | self.brick_name, self.layer_idx, self.actual_us, self.budget_us, self.gap_factor |
1532 | | ) |
1533 | 0 | } |
1534 | | } |
1535 | | |
1536 | | // ============================================================================ |
1537 | | // Benchmark Brick |
1538 | | // ============================================================================ |
1539 | | |
1540 | | /// Configuration for benchmark runs. |
1541 | | #[derive(Debug, Clone)] |
1542 | | pub struct BenchmarkConfig { |
1543 | | /// Number of warmup iterations |
1544 | | pub warmup: usize, |
1545 | | /// Number of sample iterations |
1546 | | pub samples: usize, |
1547 | | /// Maximum allowed CV (coefficient of variation) |
1548 | | pub max_cv: f64, |
1549 | | } |
1550 | | |
1551 | | impl Default for BenchmarkConfig { |
1552 | 7 | fn default() -> Self { |
1553 | 7 | Self { |
1554 | 7 | warmup: 10, |
1555 | 7 | samples: 100, |
1556 | 7 | max_cv: 0.05, // 5% per Stabilizer (Curtsinger & Berger 2013) |
1557 | 7 | } |
1558 | 7 | } |
1559 | | } |
1560 | | |
1561 | | /// Benchmark report with statistical analysis. |
1562 | | #[derive(Debug, Clone)] |
1563 | | pub struct BenchmarkReport { |
1564 | | /// Brick name |
1565 | | pub brick_name: String, |
1566 | | /// Mean latency (µs) |
1567 | | pub mean_us: f64, |
1568 | | /// Standard deviation (µs) |
1569 | | pub std_us: f64, |
1570 | | /// Coefficient of variation |
1571 | | pub cv: f64, |
1572 | | /// 50th percentile (µs) |
1573 | | pub p50_us: f64, |
1574 | | /// 99th percentile (µs) |
1575 | | pub p99_us: f64, |
1576 | | /// Throughput (tokens/sec) |
1577 | | pub tokens_per_sec: f64, |
1578 | | /// Budget target (µs) |
1579 | | pub budget_us: f64, |
1580 | | /// Budget met? |
1581 | | pub budget_met: bool, |
1582 | | /// Statistical validity (CV < max_cv) |
1583 | | pub statistically_valid: bool, |
1584 | | } |
1585 | | |
1586 | | impl fmt::Display for BenchmarkReport { |
1587 | 0 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
1588 | 0 | let status = if self.budget_met { "PASS" } else { "FAIL" }; |
1589 | 0 | write!( |
1590 | 0 | f, |
1591 | 0 | "{}: {:.1}µs ± {:.1}µs (CV={:.1}%) | {:.0} tok/s | budget: {} ({})", |
1592 | | self.brick_name, |
1593 | | self.mean_us, |
1594 | | self.std_us, |
1595 | 0 | self.cv * 100.0, |
1596 | | self.tokens_per_sec, |
1597 | | self.budget_us, |
1598 | | status |
1599 | | ) |
1600 | 0 | } |
1601 | | } |
1602 | | |
1603 | | /// Calculate percentile from sorted samples. |
1604 | 2 | fn percentile(samples: &[f64], p: f64) -> f64 { |
1605 | 2 | if samples.is_empty() { |
1606 | 0 | return 0.0; |
1607 | 2 | } |
1608 | 2 | let idx = ((samples.len() as f64) * p).floor() as usize; |
1609 | 2 | samples[idx.min(samples.len() - 1)] |
1610 | 2 | } |
1611 | | |
1612 | | /// Run benchmark on a brick with statistical rigor. |
1613 | 1 | pub fn benchmark_brick<B: ComputeBrick>( |
1614 | 1 | brick: &B, |
1615 | 1 | run_fn: impl Fn() -> f64, |
1616 | 1 | config: &BenchmarkConfig, |
1617 | 1 | ) -> BenchmarkReport { |
1618 | | // Warmup (Jidoka: ensure stable state) |
1619 | 5 | for _ in 0..config.warmup1 { |
1620 | 5 | let _ = run_fn(); |
1621 | 5 | } |
1622 | | |
1623 | | // Collect samples |
1624 | 1 | let mut samples: Vec<f64> = Vec::with_capacity(config.samples); |
1625 | 50 | for _ in 0..config.samples1 { |
1626 | 50 | samples.push(run_fn()); |
1627 | 50 | } |
1628 | | |
1629 | | // Sort for percentiles |
1630 | 293 | samples1 .sort_by1 (|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); |
1631 | | |
1632 | | // Statistical analysis |
1633 | 1 | let mean = samples.iter().sum::<f64>() / samples.len() as f64; |
1634 | 1 | let std = |
1635 | 50 | (samples.iter()1 .map1 (|x| (x - mean).powi(2)).sum1 ::<f64>() / samples.len() as f641 ).sqrt1 (); |
1636 | 1 | let cv = std / mean; |
1637 | | |
1638 | 1 | let budget = brick.budget(); |
1639 | | |
1640 | 1 | BenchmarkReport { |
1641 | 1 | brick_name: brick.name().to_string(), |
1642 | 1 | mean_us: mean, |
1643 | 1 | std_us: std, |
1644 | 1 | cv, |
1645 | 1 | p50_us: percentile(&samples, 0.50), |
1646 | 1 | p99_us: percentile(&samples, 0.99), |
1647 | 1 | tokens_per_sec: 1_000_000.0 / mean, |
1648 | 1 | budget_us: budget.us_per_token, |
1649 | 1 | budget_met: mean <= budget.us_per_token, |
1650 | 1 | statistically_valid: cv <= config.max_cv, |
1651 | 1 | } |
1652 | 1 | } |
1653 | | |
1654 | | // ============================================================================ |
1655 | | // CUDA Graph Brick (Section 5.2 - P0) |
1656 | | // ============================================================================ |
1657 | | |
1658 | | /// CUDA Graph Brick for eliminating kernel launch overhead. |
1659 | | /// |
1660 | | /// Per spec: docs/specifications/qwen2.5-coder-showcase-demo.md §5.2 |
1661 | | /// |
1662 | | /// Uses CUDA graph capture to reduce ~280 kernel launches to single graph replay. |
1663 | | /// Expected impact: 5.6ms overhead → 0.02ms = 280x overhead reduction. |
1664 | | /// |
1665 | | /// # Implementation |
1666 | | /// |
1667 | | /// Wraps `CudaExecutor::decode_graph` and `try_graph_capture()` from cuda.rs. |
1668 | | /// Uses indirect kernels (KvCacheScatterIndirect, RopeIndirect) for graph compatibility. |
1669 | | #[derive(Debug, Clone)] |
1670 | | pub struct CudaGraphBrick { |
1671 | | /// Number of layers captured in graph |
1672 | | pub num_layers: usize, |
1673 | | /// Hidden dimension |
1674 | | pub hidden_dim: usize, |
1675 | | /// Whether graph is currently captured |
1676 | | pub captured: bool, |
1677 | | /// Token budget (target: 10µs launch overhead vs 5600µs eager) |
1678 | | budget: TokenBudget, |
1679 | | } |
1680 | | |
1681 | | impl CudaGraphBrick { |
1682 | | /// Create new CUDA Graph brick for model configuration. |
1683 | | #[must_use] |
1684 | 4 | pub fn new(num_layers: usize, hidden_dim: usize) -> Self { |
1685 | | // Graph overhead should be < 100µs (vs ~5.6ms for 280 launches) |
1686 | 4 | let budget_us = 20.0; // Conservative: 20µs for graph replay |
1687 | 4 | Self { |
1688 | 4 | num_layers, |
1689 | 4 | hidden_dim, |
1690 | 4 | captured: false, |
1691 | 4 | budget: TokenBudget::from_latency(budget_us), |
1692 | 4 | } |
1693 | 4 | } |
1694 | | |
1695 | | /// Set custom budget. |
1696 | | #[must_use] |
1697 | 0 | pub fn with_budget(mut self, budget: TokenBudget) -> Self { |
1698 | 0 | self.budget = budget; |
1699 | 0 | self |
1700 | 0 | } |
1701 | | |
1702 | | /// Mark graph as captured. |
1703 | 1 | pub fn set_captured(&mut self, captured: bool) { |
1704 | 1 | self.captured = captured; |
1705 | 1 | } |
1706 | | |
1707 | | /// Check if graph can be used (captured and valid). |
1708 | | #[must_use] |
1709 | 2 | pub fn can_replay(&self) -> bool { |
1710 | 2 | self.captured |
1711 | 2 | } |
1712 | | |
1713 | | /// Replay the captured graph (stub - actual execution via CudaExecutor). |
1714 | 1 | pub fn replay(&self) -> Result<(), BrickError> { |
1715 | 1 | if !self.captured { |
1716 | 0 | return Err(BrickError::ComputeError( |
1717 | 0 | "CUDA graph not captured yet".to_string(), |
1718 | 0 | )); |
1719 | 1 | } |
1720 | | // Actual replay would be done via CudaExecutor::forward_graphed() |
1721 | 1 | Ok(()) |
1722 | 1 | } |
1723 | | } |
1724 | | |
1725 | | impl ComputeBrick for CudaGraphBrick { |
1726 | | type Output = (); |
1727 | | |
1728 | 1 | fn name(&self) -> &'static str { |
1729 | 1 | "cuda_graph" |
1730 | 1 | } |
1731 | | |
1732 | 2 | fn budget(&self) -> TokenBudget { |
1733 | 2 | self.budget |
1734 | 2 | } |
1735 | | |
1736 | 1 | fn assertions(&self) -> Vec<BrickAssertion> { |
1737 | 1 | vec![ |
1738 | 1 | BrickAssertion::budget_met(), |
1739 | 1 | BrickAssertion { |
1740 | 1 | name: "graph_speedup".to_string(), |
1741 | 1 | description: "Graph replay faster than eager execution".to_string(), |
1742 | 1 | kind: AssertionKind::Custom { |
1743 | 1 | check_name: "graph_speedup".to_string(), |
1744 | 1 | }, |
1745 | 1 | }, |
1746 | | ] |
1747 | 1 | } |
1748 | | |
1749 | 0 | fn can_run(&self) -> bool { |
1750 | 0 | self.num_layers > 0 && self.hidden_dim > 0 |
1751 | 0 | } |
1752 | | } |
1753 | | |
1754 | | // ============================================================================ |
1755 | | // Tests (F001-F020) |
1756 | | // ============================================================================ |
1757 | | |
1758 | | // Tests extracted to tests.rs (PMAT-802) |
1759 | | #[cfg(test)] |
1760 | | #[path = "tests.rs"] |
1761 | | mod brick_tests; |